From c9418306640f389ea5be974bd61b349bc583f658 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Thu, 20 Aug 2026 22:28:52 +1000 Subject: [PATCH] feat: gate self-host inference behind capacity-based ThrottlingException (#848) * feat: gate self-host inference behind capacity-based ThrottlingException Add a process-scoped concurrency limiter keyed by (servingAccountID, modelID) with capacity = catalog MaxConcurrency x max(ModelUnits, 1), and apply it to every self-host Converse/ConverseStream/InvokeModel/ InvokeModelWithResponseStream path. Over-capacity requests reject immediately with ThrottlingException instead of queuing inside vLLM, matching real Bedrock. Managed (Anthropic) providers are untouched. * test: mark admission_test as in-package for the unexported limiter * test: cover capacity admission gate and throttling --- spinifex/gateway/bedrock/admission.go | 139 +++++++++ spinifex/gateway/bedrock/admission_test.go | 336 +++++++++++++++++++++ spinifex/gateway/bedrock/catalog.go | 76 ++--- spinifex/gateway/bedrock/invoke.go | 44 +++ spinifex/gateway/bedrock/invoke_test.go | 81 +++++ spinifex/gateway/bedrock/provider.go | 30 +- spinifex/gateway/bedrock/provisioned.go | 46 +++ spinifex/gateway/bedrock/router_test.go | 160 ++++++++++ 8 files changed, 877 insertions(+), 35 deletions(-) create mode 100644 spinifex/gateway/bedrock/admission.go create mode 100644 spinifex/gateway/bedrock/admission_test.go diff --git a/spinifex/gateway/bedrock/admission.go b/spinifex/gateway/bedrock/admission.go new file mode 100644 index 000000000..4e5a91c6d --- /dev/null +++ b/spinifex/gateway/bedrock/admission.go @@ -0,0 +1,139 @@ +package gateway_bedrock + +import ( + "context" + "errors" + "log/slog" + "sync" + + "github.com/mulgadc/spinifex/spinifex/awserrors" +) + +// concurrencyLimiter is a process-scoped, in-memory admission gate keyed by +// an opaque string built from (servingAccountID, modelID). It tracks an +// authoritative in-flight count per key and admits up to a caller-supplied +// capacity, rejecting immediately rather than queuing. +type concurrencyLimiter struct { + mu sync.Mutex + inFlight map[string]int +} + +func newConcurrencyLimiter() *concurrencyLimiter { + return &concurrencyLimiter{inFlight: make(map[string]int)} +} + +// selfHostLimiter is the single process-wide limiter shared by every +// self-host Converse/ConverseStream/InvokeModel/InvokeModelWithResponseStream +// call. It is always on: there is no operator switch, matching Bedrock's own +// lack of a "disable throttling" knob. +var selfHostLimiter = newConcurrencyLimiter() + +// admissionKey builds the concurrencyLimiter key for a self-host inference +// request: servingAccountID is the resolved provisioned-throughput +// commitment's own account when a PT ARN was used, or the caller's own +// accountID for the shared GlobalAccountID-shorthand path — the same +// identity the endpoint resolver pins a serving VM on. +func admissionKey(servingAccountID, modelID string) string { + return servingAccountID + "\x00" + modelID +} + +// Acquire admits one in-flight request for key when under capacity, +// returning a release func the caller must invoke exactly once when the +// request completes. ok is false at capacity; release is then a safe no-op +// so a caller never needs to nil-check it before deferring. +func (l *concurrencyLimiter) Acquire(key string, capacity int) (release func(), ok bool) { + l.mu.Lock() + if l.inFlight[key] >= capacity { + l.mu.Unlock() + return func() {}, false + } + l.inFlight[key]++ + l.mu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + l.mu.Lock() + defer l.mu.Unlock() + l.inFlight[key]-- + if l.inFlight[key] <= 0 { + delete(l.inFlight, key) + } + }) + }, true +} + +// slotReleasingSource wraps a converseStreamSource so the admission slot +// acquired before opening the upstream request is released exactly once, +// when the stream closes, regardless of how many times Close is called. +type slotReleasingSource struct { + converseStreamSource + + release func() + once sync.Once +} + +var _ converseStreamSource = (*slotReleasingSource)(nil) + +func newSlotReleasingSource(inner converseStreamSource, release func()) *slotReleasingSource { + return &slotReleasingSource{converseStreamSource: inner, release: release} +} + +// Close delegates to the wrapped source first, then releases the admission +// slot exactly once so a slow/erroring inner Close still frees capacity. +func (s *slotReleasingSource) Close() error { + err := s.converseStreamSource.Close() + s.once.Do(s.release) + return err +} + +// slotReleasingInvokeSource is slotReleasingSource for the InvokeModelWith- +// ResponseStream path, which reframes over invokeStreamSource instead of +// converseStreamSource. +type slotReleasingInvokeSource struct { + invokeStreamSource + + release func() + once sync.Once +} + +var _ invokeStreamSource = (*slotReleasingInvokeSource)(nil) + +func newSlotReleasingInvokeSource(inner invokeStreamSource, release func()) *slotReleasingInvokeSource { + return &slotReleasingInvokeSource{invokeStreamSource: inner, release: release} +} + +// Close mirrors slotReleasingSource.Close for the invoke-stream contract. +func (s *slotReleasingInvokeSource) Close() error { + err := s.invokeStreamSource.Close() + s.once.Do(s.release) + return err +} + +// admitSelfHost resolves this self-host request's admission capacity — the +// catalog's MaxConcurrency times the serving account's committed ModelUnits, +// or times 1 on the shared ON_DEMAND path (servingAccountID empty, no +// commitment to read) — and tries to acquire a slot on selfHostLimiter. +// err is ThrottlingException when the endpoint is already at capacity; +// release must be called exactly once when the request completes on success. +func admitSelfHost(ctx context.Context, provisioned *ProvisionedStore, servingAccountID, modelID string, entry catalogEntry) (release func(), err error) { + units := int64(1) + if servingAccountID != "" { + units, err = committedModelUnits(ctx, provisioned, servingAccountID, modelID) + if err != nil { + return nil, err + } + if units < 1 { + units = 1 + } + } + capacity := entry.MaxConcurrency * int(units) + key := admissionKey(servingAccountID, modelID) + release, ok := selfHostLimiter.Acquire(key, capacity) + if !ok { + slog.Warn("bedrock: self-host request throttled", "error_code", awserrors.ErrorThrottlingException, + "model", modelID, "account", servingAccountID, "capacity", capacity) + return nil, errors.New(awserrors.ErrorThrottlingException) + } + return release, nil +} diff --git a/spinifex/gateway/bedrock/admission_test.go b/spinifex/gateway/bedrock/admission_test.go new file mode 100644 index 000000000..a43c98b10 --- /dev/null +++ b/spinifex/gateway/bedrock/admission_test.go @@ -0,0 +1,336 @@ +// slotReleasingSource Close bookkeeping directly, which cannot move to an +// external test package. +// +//test:in-package — drives the unexported concurrencyLimiter and +package gateway_bedrock + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// admissionFakeConverseStreamSource is a minimal converseStreamSource double +// for exercising slotReleasingSource's Close-call bookkeeping directly, +// independent of the pumpConverseStream-oriented fakeConverseStreamSource in +// converse_stream_test.go. +type admissionFakeConverseStreamSource struct { + closeCalls int + closeErr error +} + +var _ converseStreamSource = (*admissionFakeConverseStreamSource)(nil) + +func (f *admissionFakeConverseStreamSource) Next(_ context.Context) (ConverseStreamEvent, bool, error) { + return ConverseStreamEvent{}, false, nil +} + +func (f *admissionFakeConverseStreamSource) Close() error { + f.closeCalls++ + return f.closeErr +} + +// admissionFakeInvokeStreamSource is slotReleasingInvokeSource's counterpart +// to admissionFakeConverseStreamSource, independent of the +// pumpInvokeStream-oriented fakeInvokeStreamSource in invoke_stream_test.go. +type admissionFakeInvokeStreamSource struct { + closeCalls int + closeErr error +} + +var _ invokeStreamSource = (*admissionFakeInvokeStreamSource)(nil) + +func (f *admissionFakeInvokeStreamSource) Next(_ context.Context) ([]byte, bool, error) { + return nil, false, nil +} + +func (f *admissionFakeInvokeStreamSource) Close() error { + f.closeCalls++ + return f.closeErr +} + +func TestConcurrencyLimiter_AcquireAtCapacityRejects(t *testing.T) { + l := newConcurrencyLimiter() + + release1, ok := l.Acquire("k", 1) + require.True(t, ok) + + _, ok = l.Acquire("k", 1) + assert.False(t, ok, "a second concurrent acquire at capacity 1 must be rejected") + + release1() + + release2, ok := l.Acquire("k", 1) + require.True(t, ok, "after release, the next acquire must succeed") + release2() +} + +func TestConcurrencyLimiter_ReleaseIsIdempotent(t *testing.T) { + l := newConcurrencyLimiter() + + release, ok := l.Acquire("k", 1) + require.True(t, ok) + + release() + release() + release() + + // A double/triple release must not have driven the counter negative and + // admitted more than capacity as a result. + r2, ok := l.Acquire("k", 1) + require.True(t, ok) + _, ok = l.Acquire("k", 1) + assert.False(t, ok, "capacity must still be exactly 1 after the redundant releases") + r2() +} + +func TestConcurrencyLimiter_DifferentKeysAreIndependent(t *testing.T) { + l := newConcurrencyLimiter() + + releaseA, ok := l.Acquire("a", 1) + require.True(t, ok) + defer releaseA() + + _, ok = l.Acquire("b", 1) + assert.True(t, ok, "a different key must have its own independent capacity") +} + +// TestConcurrencyLimiter_CapacityMultipliesWithConcurrentAcquires table-drives +// capacity values shaped like catalog MaxConcurrency x ModelUnits products, +// asserting exactly `capacity` concurrent acquires are ever admitted. +func TestConcurrencyLimiter_CapacityMultipliesWithConcurrentAcquires(t *testing.T) { + tests := []struct { + name string + maxConcurrency int + modelUnits int + concurrent int + }{ + {"1B default x ON_DEMAND (units=1)", 256, 1, 300}, + {"3B x ON_DEMAND (units=1)", 8, 1, 20}, + {"3B x ModelUnits=2", 8, 2, 30}, + {"3B x ModelUnits=4", 8, 4, 40}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + l := newConcurrencyLimiter() + capacity := tt.maxConcurrency * tt.modelUnits + + admitted := 0 + for i := 0; i < tt.concurrent; i++ { + if _, ok := l.Acquire("k", capacity); ok { + admitted++ + } + } + assert.Equal(t, capacity, admitted) + }) + } +} + +func TestSlotReleasingSource_ClosesReleaseExactlyOnceAndIsIdempotent(t *testing.T) { + var mu sync.Mutex + releases := 0 + release := func() { + mu.Lock() + defer mu.Unlock() + releases++ + } + + inner := &admissionFakeConverseStreamSource{} + src := newSlotReleasingSource(inner, release) + + require.NoError(t, src.Close()) + require.NoError(t, src.Close()) + require.NoError(t, src.Close()) + + assert.Equal(t, 1, releases, "the admission slot must release exactly once") + assert.Equal(t, 3, inner.closeCalls, "every Close call must still reach the wrapped source") +} + +func TestSlotReleasingSource_ReleasesEvenWhenInnerCloseErrors(t *testing.T) { + released := false + release := func() { released = true } + + inner := &admissionFakeConverseStreamSource{closeErr: errors.New("upstream close failed")} + src := newSlotReleasingSource(inner, release) + + err := src.Close() + assert.Error(t, err) + assert.True(t, released, "the slot must release even when the inner Close errors") +} + +func TestSlotReleasingInvokeSource_ClosesReleaseExactlyOnceAndIsIdempotent(t *testing.T) { + var mu sync.Mutex + releases := 0 + release := func() { + mu.Lock() + defer mu.Unlock() + releases++ + } + + inner := &admissionFakeInvokeStreamSource{} + src := newSlotReleasingInvokeSource(inner, release) + + require.NoError(t, src.Close()) + require.NoError(t, src.Close()) + + assert.Equal(t, 1, releases) + assert.Equal(t, 2, inner.closeCalls) +} + +// TestAdmitSelfHost_OnDemandAdmitsThenThrottles drives the ON_DEMAND path +// (empty servingAccountID, units default to 1) end to end: the first request +// is admitted, the second at capacity returns a ThrottlingException, and +// capacity frees up once the in-flight release runs. +func TestAdmitSelfHost_OnDemandAdmitsThenThrottles(t *testing.T) { + ctx := context.Background() + entry := catalogEntry{MaxConcurrency: 1} + const model = "admit-ondemand-throttle-model" + + release, err := admitSelfHost(ctx, nil, "", model, entry) + require.NoError(t, err) + require.NotNil(t, release) + + _, err = admitSelfHost(ctx, nil, "", model, entry) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error(), + "a self-host request at capacity must return a ThrottlingException") + + release() + release2, err := admitSelfHost(ctx, nil, "", model, entry) + require.NoError(t, err, "capacity must free up after the in-flight request releases") + release2() +} + +// TestAdmitSelfHost_CommittedUnitsErrorPropagates asserts a store failure while +// resolving committed ModelUnits surfaces as that error, never masquerading as +// a throttle or admitting the request. +func TestAdmitSelfHost_CommittedUnitsErrorPropagates(t *testing.T) { + store := NewProvisionedStore(nil, 1, ptTestRegion, newStubEndpointProvisioner()) + + release, err := admitSelfHost(context.Background(), store, ptCallerAccount, selfHostTestModel, + catalogEntry{MaxConcurrency: 1}) + require.Error(t, err) + assert.Nil(t, release) + assert.NotEqual(t, awserrors.ErrorThrottlingException, err.Error(), + "a store failure must not masquerade as a throttle") +} + +// TestAdmitSelfHost_ClampsZeroCommittedUnits covers a serving account with no +// commitment for the model: committed units read zero and capacity clamps to +// MaxConcurrency x 1 rather than collapsing to zero and rejecting everything. +func TestAdmitSelfHost_ClampsZeroCommittedUnits(t *testing.T) { + stub := newStubEndpointProvisioner() + store := newProvisionedTestStore(t, stub) + ctx := context.Background() + entry := catalogEntry{MaxConcurrency: 1} + + release, err := admitSelfHost(ctx, store, ptCallerAccount, selfHostTestModel, entry) + require.NoError(t, err) + require.NotNil(t, release) + defer release() + + _, err = admitSelfHost(ctx, store, ptCallerAccount, selfHostTestModel, entry) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error(), + "clamped capacity is exactly 1, so a second concurrent request throttles") +} + +// TestAdmitSelfHost_StacksCommittedModelUnits proves two commitments on the +// same account+model stack their ModelUnits into one shared capacity: exactly +// MaxConcurrency x (2+3) concurrent requests are admitted before throttling. +func TestAdmitSelfHost_StacksCommittedModelUnits(t *testing.T) { + stub := newStubEndpointProvisioner() + store := newProvisionedTestStore(t, stub) + ctx := context.Background() + + _, err := CreateProvisionedModelThroughput(ctx, ptCallerAccount, store, createInput(selfHostTestModel, "pt-a", 2)) + require.NoError(t, err) + _, err = CreateProvisionedModelThroughput(ctx, ptCallerAccount, store, createInput(selfHostTestModel, "pt-b", 3)) + require.NoError(t, err) + + entry := catalogEntry{MaxConcurrency: 1} + const wantCapacity = 5 + + releases := make([]func(), 0, wantCapacity) + for i := range wantCapacity { + release, err := admitSelfHost(ctx, store, ptCallerAccount, selfHostTestModel, entry) + require.NoError(t, err, "acquire %d must be admitted within stacked capacity", i) + releases = append(releases, release) + } + _, err = admitSelfHost(ctx, store, ptCallerAccount, selfHostTestModel, entry) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error(), + "the request past stacked capacity must throttle") + + for _, release := range releases { + release() + } +} + +// TestCommittedModelUnits_EmptyAccountReturnsZero covers the shared ON_DEMAND +// short-circuit: an empty serving account has no commitment to read and must +// not touch the store, so a nil store is safe. +func TestCommittedModelUnits_EmptyAccountReturnsZero(t *testing.T) { + units, err := committedModelUnits(context.Background(), nil, "", selfHostTestModel) + require.NoError(t, err) + assert.Equal(t, int64(0), units) +} + +// TestCommittedModelUnits_BucketErrorPropagates asserts an unusable store +// surfaces its error rather than silently reading zero committed units. +func TestCommittedModelUnits_BucketErrorPropagates(t *testing.T) { + store := NewProvisionedStore(nil, 1, ptTestRegion, newStubEndpointProvisioner()) + _, err := committedModelUnits(context.Background(), store, ptCallerAccount, selfHostTestModel) + require.Error(t, err) +} + +// TestCommittedModelUnits_NoCommitmentsReturnsZero covers the empty-bucket path: +// a store with no commitments at all reads zero without erroring. +func TestCommittedModelUnits_NoCommitmentsReturnsZero(t *testing.T) { + store := newProvisionedTestStore(t, newStubEndpointProvisioner()) + units, err := committedModelUnits(context.Background(), store, ptCallerAccount, selfHostTestModel) + require.NoError(t, err) + assert.Equal(t, int64(0), units) +} + +// TestCommittedModelUnits_ScopesToAccountAndModel proves the sum stacks only a +// caller's own commitments for the requested model, skipping a different model +// and a different account, and reads zero for a model with no commitment. +func TestCommittedModelUnits_ScopesToAccountAndModel(t *testing.T) { + store := newProvisionedTestStore(t, newStubEndpointProvisioner()) + ctx := context.Background() + + mustCreate := func(account, model, name string, units int64) { + t.Helper() + _, err := CreateProvisionedModelThroughput(ctx, account, store, createInput(model, name, units)) + require.NoError(t, err) + } + mustCreate(ptCallerAccount, selfHostTestModel, "caller-a", 2) + mustCreate(ptCallerAccount, selfHostTestModel, "caller-b", 3) + mustCreate(ptCallerAccount, selfHostTestModel3B, "other-model", 4) + mustCreate(ptOtherCaller, selfHostTestModel, "other-account", 9) + + units, err := committedModelUnits(ctx, store, ptCallerAccount, selfHostTestModel) + require.NoError(t, err) + assert.Equal(t, int64(5), units, "only the caller's own commitments for this model stack") + + unknown, err := committedModelUnits(ctx, store, ptCallerAccount, "meta.does-not-exist-v1:0") + require.NoError(t, err) + assert.Equal(t, int64(0), unknown, "a model with no commitment reads zero") +} + +func TestAdmissionKey_ScopesByAccountAndModel(t *testing.T) { + assert.NotEqual(t, admissionKey("a", "m"), admissionKey("b", "m"), "different accounts must not share a key") + assert.NotEqual(t, admissionKey("a", "m1"), admissionKey("a", "m2"), "different models must not share a key") + + l := newConcurrencyLimiter() + release, ok := l.Acquire(admissionKey("a", "m"), 1) + require.True(t, ok) + defer release() + + _, ok = l.Acquire(admissionKey("a", "m"), 1) + assert.False(t, ok, "the same (account, model) pair must produce the same key and share capacity") +} diff --git a/spinifex/gateway/bedrock/catalog.go b/spinifex/gateway/bedrock/catalog.go index dad2c3048..e27d04600 100644 --- a/spinifex/gateway/bedrock/catalog.go +++ b/spinifex/gateway/bedrock/catalog.go @@ -54,25 +54,29 @@ const ( // empty value is a standalone model (a bundle of one), a shared non-empty // value means the entries admit and launch together as a single GPU claim. type catalogEntry struct { - ModelID string - ModelName string - ProviderName string - Provider string // "self-host" or "provider:" - Family string // self-host only; selects the invoke adapter (see familyMeta, familyTEI) - CoServeGroup string // self-host only; shared co-serve group id, empty = standalone - HFRepo string // self-host only; canonical Hugging Face repo 'weights pull' defaults to - HFRevision string // self-host only; pinned HF commit/tag 'weights pull' defaults to, empty = main - InputModalities []string - OutputModalities []string - ResponseStreamingSupported bool - InferenceTypesSupported []string - CustomizationsSupported []string - MinVRAMMiB int // self-host only; 0 for provider entries - InstanceType string // self-host only; system-instance type a launcher boots - VLLMArgs []string // self-host only; extra vLLM server args this model needs - InputPriceMicroUSDPerMillion int64 // provider entries only; meaningless unless PriceKnown - OutputPriceMicroUSDPerMillion int64 // provider entries only; meaningless unless PriceKnown - PriceKnown bool // provider entries only; self-host is always known-zero + ModelID string + ModelName string + ProviderName string + Provider string // "self-host" or "provider:" + Family string // self-host only; selects the invoke adapter (see familyMeta, familyTEI) + CoServeGroup string // self-host only; shared co-serve group id, empty = standalone + HFRepo string // self-host only; canonical Hugging Face repo 'weights pull' defaults to + HFRevision string // self-host only; pinned HF commit/tag 'weights pull' defaults to, empty = main + InputModalities []string + OutputModalities []string + ResponseStreamingSupported bool + InferenceTypesSupported []string + CustomizationsSupported []string + MinVRAMMiB int // self-host only; 0 for provider entries + InstanceType string // self-host only; system-instance type a launcher boots + VLLMArgs []string // self-host only; extra vLLM server args this model needs + // MaxConcurrency is the self-host admission-gate capacity and must stay + // equal to the entry's --max-num-seqs VLLMArgs value (or vLLM's own + // default when that flag is absent) — edit the two together. + MaxConcurrency int + InputPriceMicroUSDPerMillion int64 // provider entries only; meaningless unless PriceKnown + OutputPriceMicroUSDPerMillion int64 // provider entries only; meaningless unless PriceKnown + PriceKnown bool // provider entries only; self-host is always known-zero } // catalog is the static model set: two self-hosted open models and one @@ -93,9 +97,10 @@ var catalog = []catalogEntry{ // MinVRAMMiB is the admission-gate floor; --gpu-memory-utilization // caps vLLM's own pool to roughly the same figure (8188 MiB * 0.6 ≈ // 4913 MiB) so the two stay consistent rather than drifting apart. - MinVRAMMiB: 5120, - InstanceType: "g5.xlarge", - VLLMArgs: []string{"--dtype=bfloat16", "--max-model-len=8192", "--gpu-memory-utilization=0.6"}, + MinVRAMMiB: 5120, + InstanceType: "g5.xlarge", + VLLMArgs: []string{"--dtype=bfloat16", "--max-model-len=8192", "--gpu-memory-utilization=0.6"}, + MaxConcurrency: 256, }, { ModelID: "meta.llama3-2-3b-instruct-v1:0", @@ -117,6 +122,7 @@ var catalog = []catalogEntry{ "--dtype=bfloat16", "--max-model-len=4096", "--gpu-memory-utilization=0.92", "--max-num-seqs=8", "--enforce-eager", }, + MaxConcurrency: 8, }, { ModelID: "nomic-embed-text-v1.5", @@ -431,12 +437,13 @@ func LookupCoServeGroup(modelID string) (spec CoServeGroupSpec, found, selfHost // before staging weights) and by the daemon-side launcher (placing and // booting the serving VM), without exposing catalogEntry's AWS-shaped fields. type ServingSpec struct { - ModelID string - MinVRAMMiB int - InstanceType string - VLLMArgs []string - HFRepo string - HFRevision string + ModelID string + MinVRAMMiB int + InstanceType string + VLLMArgs []string + HFRepo string + HFRevision string + MaxConcurrency int } // LookupServingSpec returns modelID's serving spec. found reports whether @@ -452,12 +459,13 @@ func LookupServingSpec(modelID string) (spec ServingSpec, found, selfHost bool) return ServingSpec{}, true, false } return ServingSpec{ - ModelID: entry.ModelID, - MinVRAMMiB: entry.MinVRAMMiB, - InstanceType: entry.InstanceType, - VLLMArgs: entry.VLLMArgs, - HFRepo: entry.HFRepo, - HFRevision: entry.HFRevision, + ModelID: entry.ModelID, + MinVRAMMiB: entry.MinVRAMMiB, + InstanceType: entry.InstanceType, + VLLMArgs: entry.VLLMArgs, + HFRepo: entry.HFRepo, + HFRevision: entry.HFRevision, + MaxConcurrency: entry.MaxConcurrency, }, true, true } diff --git a/spinifex/gateway/bedrock/invoke.go b/spinifex/gateway/bedrock/invoke.go index 95e401168..91490479f 100644 --- a/spinifex/gateway/bedrock/invoke.go +++ b/spinifex/gateway/bedrock/invoke.go @@ -133,10 +133,19 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri var a InvokeAdapter switch { case entry.Provider == tierSelfHost: + // Family validation first: a model this platform cannot serve at + // all must fail ValidationException outright, not consume (or be + // refused for) admission capacity it was never going to use. a, err = selfHostInvokeAdapter(entry.Family, rt.endpointResolver, ptAccountID) if err != nil { return nil, "", err } + var release func() + release, err = admitSelfHost(ctx, rt.provisioned, ptAccountID, modelID, entry) + if err != nil { + return nil, "", err + } + defer release() case strings.HasPrefix(entry.Provider, providerPrefix): switch strings.TrimPrefix(entry.Provider, providerPrefix) { case vendorAnthropic: @@ -299,12 +308,24 @@ func (rt *InvokeStreamRouter) InvokeModelWithResponseStream(ctx context.Context, } var a InvokeAdapter + // selfHostRelease is non-nil only on the self-host branch below; every + // return past this point on that branch must release the slot exactly + // once, either explicitly here or via the slot-releasing source wrapper + // on the success path. + var selfHostRelease func() switch { case entry.Provider == tierSelfHost: + // Family validation first: a model this platform cannot serve at + // all must fail ValidationException outright, not consume (or be + // refused for) admission capacity it was never going to use. a, err = selfHostInvokeAdapter(entry.Family, rt.endpointResolver, ptAccountID) if err != nil { return nil, err } + selfHostRelease, err = admitSelfHost(ctx, rt.provisioned, ptAccountID, modelID, entry) + if err != nil { + return nil, err + } case strings.HasPrefix(entry.Provider, providerPrefix): switch strings.TrimPrefix(entry.Provider, providerPrefix) { case vendorAnthropic: @@ -325,32 +346,55 @@ func (rt *InvokeStreamRouter) InvokeModelWithResponseStream(ctx context.Context, sa, ok := a.(InvokeStreamAdapter) if !ok { + if selfHostRelease != nil { + selfHostRelease() + } return nil, errors.New(awserrors.ErrorValidationException) } if guardrailIdent != "" { texts, extractOK := extractInvokePromptTexts(entry.Provider, body) if !extractOK { + if selfHostRelease != nil { + selfHostRelease() + } return nil, errors.New(awserrors.ErrorValidationException) } blocked, message, _, _, gerr := enforceGuardrail(ctx, rt.guardrails, accountID, guardrailIdent, guardrailVersion, bedrockruntime.GuardrailContentSourceInput, texts) if gerr != nil { + if selfHostRelease != nil { + selfHostRelease() + } return nil, gerr } if blocked { chunks, berr := buildGuardedInvokeStreamChunks(entry.Provider, message, true) if berr != nil { + if selfHostRelease != nil { + selfHostRelease() + } return nil, berr } + // A blocked INPUT never opens the backend stream, so the slot + // releases immediately rather than living past this return. + if selfHostRelease != nil { + selfHostRelease() + } return &blockedInvokeStreamSource{chunks: chunks}, nil } } src, err := sa.InvokeModelWithResponseStream(ctx, modelID, body) if err != nil { + if selfHostRelease != nil { + selfHostRelease() + } return nil, err } + if selfHostRelease != nil { + src = newSlotReleasingInvokeSource(src, selfHostRelease) + } if guardrailIdent != "" { src = newGuardrailInvokeStreamSource(src, rt.guardrails, accountID, guardrailIdent, guardrailVersion, entry.Provider) } diff --git a/spinifex/gateway/bedrock/invoke_test.go b/spinifex/gateway/bedrock/invoke_test.go index cbae8d0c7..163b85fbf 100644 --- a/spinifex/gateway/bedrock/invoke_test.go +++ b/spinifex/gateway/bedrock/invoke_test.go @@ -192,3 +192,84 @@ func TestInvokeStreamRouter_SelfHostUnhandledFamilyReturnsValidationException(t assert.Equal(t, awserrors.ErrorValidationException, err.Error()) assert.False(t, called, "an unhandled self-host family must never reach the Llama-serving endpoint") } + +// TestInvokeRouter_SelfHostGatedAtCapacity proves InvokeModel's self-host +// path shares the same admission gate as Converse: pre-occupying the +// endpoint's only slot throttles the next call, and releasing it admits +// the following one. +func TestInvokeRouter_SelfHostGatedAtCapacity(t *testing.T) { + modelID := "self-host.throttle-invoke-v1:0" + withCatalogEntry(t, catalogEntry{ModelID: modelID, Provider: tierSelfHost, Family: familyMeta, MaxConcurrency: 1}) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "choices": [{"text": "hi", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil) + + release, ok := selfHostLimiter.Acquire(admissionKey("", modelID), 1) + require.True(t, ok) + + _, _, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) + + release() + + respBody, _, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "") + require.NoError(t, err) + + var out llamaInvokeResponse + require.NoError(t, json.Unmarshal(respBody, &out)) + assert.Equal(t, "hi", out.Generation) +} + +// TestInvokeStreamRouter_SelfHostGatedAtCapacity mirrors the non-stream case +// for InvokeModelWithResponseStream, and proves the slot stays held for the +// whole stream lifetime — released only on the returned source's Close. +func TestInvokeStreamRouter_SelfHostGatedAtCapacity(t *testing.T) { + modelID := "self-host.throttle-invoke-stream-v1:0" + withCatalogEntry(t, catalogEntry{ModelID: modelID, Provider: tierSelfHost, Family: familyMeta, MaxConcurrency: 1}) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(llamaCompletionsStreamFixture)) + })) + defer ts.Close() + + rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), grantAll{}, nil, nil) + + release, ok := selfHostLimiter.Acquire(admissionKey("", modelID), 1) + require.True(t, ok) + + _, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) + + release() + + src, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "") + require.NoError(t, err) + chunks := drainInvokeStream(t, src) + assert.NotEmpty(t, chunks) + + // The stream is drained but not yet closed: its own acquire still holds + // the endpoint's only slot, so a further concurrent request must throttle. + _, err = rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) + + require.NoError(t, src.Close()) + + // Close released the slot; the endpoint is admissible again. + src2, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), "", "") + require.NoError(t, err) + require.NoError(t, src2.Close()) +} diff --git a/spinifex/gateway/bedrock/provider.go b/spinifex/gateway/bedrock/provider.go index 6005c3e47..e453e9989 100644 --- a/spinifex/gateway/bedrock/provider.go +++ b/spinifex/gateway/bedrock/provider.go @@ -131,6 +131,12 @@ func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input var p Provider switch { case entry.Provider == tierSelfHost: + var release func() + release, err = admitSelfHost(ctx, rt.provisioned, ptAccountID, modelID, entry) + if err != nil { + return nil, err + } + defer release() if ptAccountID != "" { p = newVLLMProviderForAccount(rt.endpointResolver, ptAccountID) } else { @@ -253,8 +259,16 @@ func (rt *Router) ConverseStream(ctx context.Context, accountID, modelID string, } var p Provider + // selfHostRelease is non-nil only on the self-host branch below; it gates + // whether the eventual source gets slot-release wrapping and whether an + // error return past this point must release the slot explicitly. + var selfHostRelease func() switch { case entry.Provider == tierSelfHost: + selfHostRelease, err = admitSelfHost(ctx, rt.provisioned, ptAccountID, modelID, entry) + if err != nil { + return nil, err + } if ptAccountID != "" { p = newVLLMProviderForAccount(rt.endpointResolver, ptAccountID) } else { @@ -280,7 +294,21 @@ func (rt *Router) ConverseStream(ctx context.Context, accountID, modelID string, sp, ok := p.(ConverseStreamProvider) if !ok { + if selfHostRelease != nil { + selfHostRelease() + } return nil, errors.New(awserrors.ErrorValidationException) } - return sp.ConverseStream(ctx, modelID, input) + + src, err := sp.ConverseStream(ctx, modelID, input) + if err != nil { + if selfHostRelease != nil { + selfHostRelease() + } + return nil, err + } + if selfHostRelease != nil { + src = newSlotReleasingSource(src, selfHostRelease) + } + return src, nil } diff --git a/spinifex/gateway/bedrock/provisioned.go b/spinifex/gateway/bedrock/provisioned.go index a515cd875..0cf8741c7 100644 --- a/spinifex/gateway/bedrock/provisioned.go +++ b/spinifex/gateway/bedrock/provisioned.go @@ -146,6 +146,52 @@ func getProvisionedRecordRevision(ctx context.Context, kv jetstream.KeyValue, ke return rec, entry.Revision(), true, nil } +// committedModelUnits sums ModelUnits across every commitment servingAccountID +// holds for modelID. Every such commitment pins the same daemon endpoint +// (EnsurePinned is keyed on (accountID, modelID) alone), so their capacity +// stacks onto one shared serving VM rather than being read independently. +// An empty servingAccountID (the shared ON_DEMAND path, which has no +// commitment to read) returns 0 without touching the store. +func committedModelUnits(ctx context.Context, store *ProvisionedStore, servingAccountID, modelID string) (int64, error) { + if servingAccountID == "" { + return 0, nil + } + kv, err := store.bucket(ctx) + if err != nil { + return 0, err + } + keys, err := kv.Keys(ctx) + if err != nil { + if errors.Is(err, jetstream.ErrNoKeysFound) { + return 0, nil + } + return 0, fmt.Errorf("kv keys for provisioned models: %w", err) + } + + prefix := servingAccountID + "/" + var total int64 + for _, key := range keys { + if !strings.HasPrefix(key, prefix) { + continue + } + entry, err := kv.Get(ctx, key) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + continue + } + return 0, fmt.Errorf("kv get %s: %w", key, err) + } + var rec ProvisionedModelRecord + if err := json.Unmarshal(entry.Value(), &rec); err != nil { + continue + } + if rec.ModelID == modelID { + total += rec.ModelUnits + } + } + return total, nil +} + // deriveStatus reports the AWS-shaped status for (accountID, modelID)'s // pinned endpoint. Anything other than STARTING/READY — including an absent // endpoint, DRAINING (only ever transient during this package's own Delete), diff --git a/spinifex/gateway/bedrock/router_test.go b/spinifex/gateway/bedrock/router_test.go index 3bfa2aac8..96d1c9544 100644 --- a/spinifex/gateway/bedrock/router_test.go +++ b/spinifex/gateway/bedrock/router_test.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "net/http/httptest" + "sync" "testing" "github.com/aws/aws-sdk-go/aws" @@ -150,3 +151,162 @@ func TestNewRouter_NilArgumentsFallBackToNoops(t *testing.T) { require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } + +// TestRouter_Converse_SelfHostGatedAtCapacity proves the non-stream self-host +// path is admission-gated: pre-occupying the endpoint's only slot throttles +// the next Converse, and releasing it lets the following one through. +func TestRouter_Converse_SelfHostGatedAtCapacity(t *testing.T) { + modelID := "self-host.throttle-converse-v1:0" + withCatalogEntry(t, catalogEntry{ModelID: modelID, Provider: tierSelfHost, Family: familyMeta, MaxConcurrency: 1}) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil) + + // Occupy the endpoint's only admission slot directly, standing in for a + // concurrent in-flight request on the same (account, model) key. + release, ok := selfHostLimiter.Acquire(admissionKey("", modelID), 1) + require.True(t, ok) + + _, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput()) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) + + release() + + out, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput()) + require.NoError(t, err) + require.NotNil(t, out.Output.Message) +} + +// TestRouter_ConverseStream_SelfHostGatedAtCapacity mirrors the non-stream +// case for ConverseStream, and additionally proves the slot stays held for +// the whole stream lifetime — released only on the returned source's Close. +func TestRouter_ConverseStream_SelfHostGatedAtCapacity(t *testing.T) { + modelID := "self-host.throttle-converse-stream-v1:0" + withCatalogEntry(t, catalogEntry{ModelID: modelID, Provider: tierSelfHost, Family: familyMeta, MaxConcurrency: 1}) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(vllmStreamFixture)) + })) + defer ts.Close() + + rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, nil) + + release, ok := selfHostLimiter.Acquire(admissionKey("", modelID), 1) + require.True(t, ok) + + _, err := rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput()) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) + + release() + + src, err := rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput()) + require.NoError(t, err) + events := drainConverseStream(t, src) + assert.Equal(t, converseStreamEventMessageStart, events[0].Kind) + + // The stream is drained but not yet closed: its own acquire still holds + // the endpoint's only slot, so a further concurrent request must throttle. + _, err = rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput()) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) + + require.NoError(t, src.Close()) + + // Close released the slot; the endpoint is admissible again. + src2, err := rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput()) + require.NoError(t, err) + require.NoError(t, src2.Close()) +} + +// TestRouter_Converse_AnthropicPathIsNeverGated fires more concurrent +// Anthropic Converse calls than any self-host capacity in the catalog and +// asserts every one fails on the credential check (AccessDeniedException), +// never ThrottlingException — the managed-provider branch must never reach +// admitSelfHost. +func TestRouter_Converse_AnthropicPathIsNeverGated(t *testing.T) { + rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil, grantAll{}, nil, nil) + + const n = 20 + errs := make([]error, n) + var wg sync.WaitGroup + for i := range n { + wg.Add(1) + go func(i int) { + defer wg.Done() + _, errs[i] = rt.Converse(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", converseInput()) + }(i) + } + wg.Wait() + + for _, err := range errs { + require.Error(t, err) + assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error()) + } +} + +// TestRouter_Converse_ProvisionedThroughputARN_CapacityMultipliesByModelUnits +// proves D2's capacity formula end to end: a PT ARN commitment with +// ModelUnits=3 against a MaxConcurrency=1 catalog entry admits exactly 3 +// concurrent requests, not 1. +func TestRouter_Converse_ProvisionedThroughputARN_CapacityMultipliesByModelUnits(t *testing.T) { + modelID := "self-host.throttle-pt-model-units-v1:0" + withCatalogEntry(t, catalogEntry{ModelID: modelID, Provider: tierSelfHost, Family: familyMeta, MaxConcurrency: 1}) + + store := newProvisionedTestStore(t, newStubEndpointProvisioner()) + out, err := CreateProvisionedModelThroughput(context.Background(), ptCallerAccount, store, createInput(modelID, "my-pt", 3)) + require.NoError(t, err) + arn := aws.StringValue(out.ProvisionedModelArn) + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, store, nil) + key := admissionKey(ptCallerAccount, modelID) + + // Occupy 1 of the 3 units directly (matching MaxConcurrency alone). If + // the router only consulted MaxConcurrency and ignored ModelUnits, this + // single occupied slot would already read as "at capacity" and the + // Converse below would incorrectly throttle. + release1, ok := selfHostLimiter.Acquire(key, 3) + require.True(t, ok) + + out2, err := rt.Converse(context.Background(), ptCallerAccount, arn, converseInput()) + require.NoError(t, err, "capacity must be MaxConcurrency x ModelUnits (3), not MaxConcurrency alone (1)") + require.NotNil(t, out2.Output.Message) + + // Occupy the remaining 2 units directly to reach the 3-unit ceiling — + // Converse's own transient acquire/release above already gave back its + // slot, so inFlight is back to 1 (just release1) at this point. + release2, ok := selfHostLimiter.Acquire(key, 3) + require.True(t, ok) + release3, ok := selfHostLimiter.Acquire(key, 3) + require.True(t, ok) + + _, err = rt.Converse(context.Background(), ptCallerAccount, arn, converseInput()) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorThrottlingException, err.Error()) + + release1() + release2() + release3() +}