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
2026-10-15 17:11:30 +00:00 · 2026-08-20 22:28:52 +10:00
parent c918f6834a
commit c941830664
8 changed files with 877 additions and 35 deletions
@@ -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
+}
@@ -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")
+}
@@ -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:<vendor>"
-	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:<vendor>"
+	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
 }

@@ -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)
 	}
@@ -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())
+}
@@ -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
 }
@@ -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),
@@ -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()
+}