diff --git a/spinifex/gateway/bedrock.go b/spinifex/gateway/bedrock.go index 14c26f890..eeea1440a 100644 --- a/spinifex/gateway/bedrock.go +++ b/spinifex/gateway/bedrock.go @@ -173,9 +173,13 @@ func (gw *GatewayConfig) bedrockAccessResolver() gateway_bedrock.AccessResolver return gateway_bedrock.DenyAllAccessResolver } -// bedrockEndpointResolver returns an EndpointResolver over the configured -// pinned self-host endpoints (gw.BedrockEndpoints). A nil/empty map resolves -// nothing, so self-host models return ModelNotReady until configured. +// bedrockEndpointResolver returns the registry-backed resolver when one is +// configured, which resolves gw.BedrockEndpoints ahead of the registry itself. +// Without it only the pinned endpoints resolve, so a model that was never +// pinned returns ModelNotReady rather than launching. func (gw *GatewayConfig) bedrockEndpointResolver() gateway_bedrock.EndpointResolver { + if gw.BedrockEndpointResolver != nil { + return gw.BedrockEndpointResolver + } return gateway_bedrock.NewStaticEndpointResolver(gw.BedrockEndpoints) } diff --git a/spinifex/gateway/bedrock_test.go b/spinifex/gateway/bedrock_test.go index 9a52cc718..1f0812712 100644 --- a/spinifex/gateway/bedrock_test.go +++ b/spinifex/gateway/bedrock_test.go @@ -1,8 +1,11 @@ package gateway import ( + "context" "testing" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -46,3 +49,37 @@ func TestLookupBedrockRuntimeAction_UnknownReturnsFalse(t *testing.T) { assert.False(t, ok) assert.Nil(t, handler) } + +// stubEndpointResolver stands in for the registry-backed resolver, so the +// preference test needs no NATS connection. +type stubEndpointResolver struct{ baseURL string } + +var _ gateway_bedrock.EndpointResolver = stubEndpointResolver{} + +func (s stubEndpointResolver) Endpoint(context.Context, string) (string, bool, error) { + return s.baseURL, s.baseURL != "", nil +} + +// TestBedrockEndpointResolver_PrefersDynamic pins the wiring the whole +// dynamic-resolution change hangs off: with a registry resolver configured the +// gateway must use it, since it is the only path that can request a launch. +func TestBedrockEndpointResolver_PrefersDynamic(t *testing.T) { + gw := &GatewayConfig{ + BedrockEndpoints: map[string]string{"m": "http://static:8000"}, + BedrockEndpointResolver: stubEndpointResolver{baseURL: "http://dynamic:8000"}, + } + baseURL, ok, err := gw.bedrockEndpointResolver().Endpoint(context.Background(), "m") + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "http://dynamic:8000", baseURL) +} + +// TestBedrockEndpointResolver_FallsBackToStatic keeps every gateway built +// without a registry resolver (including the tests) on its pinned map. +func TestBedrockEndpointResolver_FallsBackToStatic(t *testing.T) { + gw := &GatewayConfig{BedrockEndpoints: map[string]string{"m": "http://static:8000"}} + baseURL, ok, err := gw.bedrockEndpointResolver().Endpoint(context.Background(), "m") + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "http://static:8000", baseURL) +} diff --git a/spinifex/gateway/gateway.go b/spinifex/gateway/gateway.go index a06182691..11356dd4d 100644 --- a/spinifex/gateway/gateway.go +++ b/spinifex/gateway/gateway.go @@ -121,9 +121,14 @@ type GatewayConfig struct { BedrockCredentials *gateway_bedrock.CredentialStore // BedrockEndpoints maps a self-hosted modelId to its OpenAI-compatible base // URL. These endpoints are pinned/always-resident, so this is static - // config; a future implementation can back it with the daemon's dynamic endpoint - // registry. Nil/empty means no self-hosted models are reachable. + // config, and they resolve ahead of the dynamic registry so a pinned + // endpoint bypasses the lifecycle entirely. BedrockEndpoints map[string]string + // BedrockEndpointResolver resolves self-host endpoints through the daemon's + // registry, requesting a launch for a model that has no endpoint yet. Nil + // falls back to BedrockEndpoints alone, under which a model that was never + // pinned is never reachable. + BedrockEndpointResolver gateway_bedrock.EndpointResolver // BedrockLoggingConfig persists per-account invocation-logging preferences // (PutModelInvocationLoggingConfiguration and friends). Nil falls back to // an unconfigured store, under which reads/writes error rather than panic. diff --git a/spinifex/handlers/bedrock/resolver.go b/spinifex/handlers/bedrock/resolver.go new file mode 100644 index 000000000..93c6118e3 --- /dev/null +++ b/spinifex/handlers/bedrock/resolver.go @@ -0,0 +1,126 @@ +package handlers_bedrock + +import ( + "context" + "sync" + "time" + + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" + "github.com/mulgadc/spinifex/spinifex/utils" + "golang.org/x/sync/singleflight" +) + +// defaultEndpointCacheTTL bounds how long a READY base URL is reused without +// re-describing. Short, because the record can go away underneath the cache +// (a delete, or idle reclaim later), and the cost of being wrong is calls to +// an address that no longer answers. +const defaultEndpointCacheTTL = 5 * time.Second + +// 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 +// cannot import this one: handlers_bedrock already imports it for +// LookupServingSpec. +type DynamicEndpointResolver struct { + svc EndpointService + static map[string]string + ttl 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 + // trip per request to learn the same answer. + group singleflight.Group + + mu sync.Mutex + cached map[string]cachedEndpoint +} + +// cachedEndpoint is one resolved READY base URL and the moment it expires. +type cachedEndpoint struct { + baseURL string + expiresAt time.Time +} + +var _ gateway_bedrock.EndpointResolver = (*DynamicEndpointResolver)(nil) + +// 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 { + if ttl <= 0 { + ttl = defaultEndpointCacheTTL + } + return &DynamicEndpointResolver{ + svc: svc, + static: static, + ttl: ttl, + cached: make(map[string]cachedEndpoint), + } +} + +// Endpoint returns modelID's base URL if one is serving. A model with no +// endpoint yet is requested from the daemon and reported as unresolved: the +// invoke paths turn that into ModelNotReadyException, so a cold call gets a +// retryable answer immediately and a retry once the VM is up succeeds. +func (r *DynamicEndpointResolver) Endpoint(ctx context.Context, modelID string) (string, bool, error) { + if baseURL, ok := r.static[modelID]; ok { + return baseURL, true, nil + } + if baseURL, ok := r.lookupCache(modelID); ok { + return baseURL, true, nil + } + + resolved, err, _ := r.group.Do(modelID, func() (any, error) { + return r.resolve(ctx, modelID) + }) + if err != nil { + return "", false, err + } + baseURL, _ := resolved.(string) + return baseURL, baseURL != "", nil +} + +// resolve describes the endpoint and, when there is none, asks for one. An +// empty base URL means "not resolved", which is the only outcome a cold model +// can have: the launch outlives this request by design. +func (r *DynamicEndpointResolver) resolve(ctx context.Context, modelID string) (string, error) { + out, err := r.svc.Describe(ctx, &DescribeEndpointInput{ModelID: modelID}, utils.GlobalAccountID) + if err != nil { + return "", err + } + + switch out.Endpoint.State { + case StateReady: + if out.Endpoint.BaseURL == "" { + return "", nil + } + 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. + return "", nil + } + + if _, err := r.svc.Ensure(ctx, &EnsureEndpointInput{ModelID: modelID}, utils.GlobalAccountID); err != nil { + return "", err + } + return "", nil +} + +func (r *DynamicEndpointResolver) lookupCache(modelID string) (string, bool) { + r.mu.Lock() + defer r.mu.Unlock() + entry, ok := r.cached[modelID] + if !ok || time.Now().After(entry.expiresAt) { + return "", false + } + return entry.baseURL, true +} + +func (r *DynamicEndpointResolver) storeCache(modelID, baseURL string) { + r.mu.Lock() + defer r.mu.Unlock() + r.cached[modelID] = cachedEndpoint{baseURL: baseURL, expiresAt: time.Now().Add(r.ttl)} +} diff --git a/spinifex/handlers/bedrock/resolver_test.go b/spinifex/handlers/bedrock/resolver_test.go new file mode 100644 index 000000000..619c8c5f1 --- /dev/null +++ b/spinifex/handlers/bedrock/resolver_test.go @@ -0,0 +1,241 @@ +package handlers_bedrock + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const resolverTestModel = "meta.llama3-2-1b-instruct-v1:0" + +// fakeEndpointService is an EndpointService returning a scripted record and +// counting the calls, so a resolver test can assert both the answer and how +// many NATS round trips producing it would have cost. +type fakeEndpointService struct { + mu sync.Mutex + + record EndpointRecord + describeErr error + ensureErr error + + describeCalls atomic.Int64 + ensureCalls atomic.Int64 + + // beforeDescribe runs inside Describe, so a test can hold every caller at + // the same point to exercise the concurrent path. + beforeDescribe func() +} + +var _ EndpointService = (*fakeEndpointService)(nil) + +func (f *fakeEndpointService) Describe(_ context.Context, _ *DescribeEndpointInput, _ string) (*DescribeEndpointOutput, error) { + f.describeCalls.Add(1) + if f.beforeDescribe != nil { + f.beforeDescribe() + } + if f.describeErr != nil { + return nil, f.describeErr + } + f.mu.Lock() + defer f.mu.Unlock() + return &DescribeEndpointOutput{Endpoint: f.record}, nil +} + +// Ensure records the request and moves the fake to STARTING, mirroring the +// daemon: the reply is immediate and the launch outlives it. +func (f *fakeEndpointService) Ensure(_ context.Context, in *EnsureEndpointInput, _ string) (*EnsureEndpointOutput, error) { + f.ensureCalls.Add(1) + if f.ensureErr != nil { + return nil, f.ensureErr + } + f.mu.Lock() + defer f.mu.Unlock() + f.record = EndpointRecord{ModelID: in.ModelID, State: StateStarting} + return &EnsureEndpointOutput{Endpoint: f.record}, nil +} + +func (f *fakeEndpointService) List(context.Context, *ListEndpointsInput, string) (*ListEndpointsOutput, error) { + return &ListEndpointsOutput{}, nil +} + +func (f *fakeEndpointService) Delete(context.Context, *DeleteEndpointInput, string) (*DeleteEndpointOutput, error) { + return &DeleteEndpointOutput{}, nil +} + +func (f *fakeEndpointService) setRecord(rec EndpointRecord) { + f.mu.Lock() + defer f.mu.Unlock() + f.record = rec +} + +func readyRecord(baseURL string) EndpointRecord { + return EndpointRecord{ModelID: resolverTestModel, State: StateReady, BaseURL: baseURL} +} + +// TestDynamicEndpointResolver_StaticWins keeps the escape hatch: a pinned +// endpoint bypasses the lifecycle entirely, which is what dev boxes and the +// gated E2E tier rely on. +func TestDynamicEndpointResolver_StaticWins(t *testing.T) { + svc := &fakeEndpointService{record: readyRecord("http://10.0.0.9:8000")} + r := NewDynamicEndpointResolver(svc, map[string]string{resolverTestModel: "http://pinned:8000"}, 0) + + baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, "http://pinned:8000", baseURL) + assert.Zero(t, svc.describeCalls.Load(), "a pinned endpoint must not consult the registry") +} + +func TestDynamicEndpointResolver_ReadyResolves(t *testing.T) { + svc := &fakeEndpointService{record: readyRecord("http://10.0.0.9:8000")} + r := NewDynamicEndpointResolver(svc, nil, 0) + + baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, "http://10.0.0.9:8000", baseURL) + assert.Zero(t, svc.ensureCalls.Load(), "a serving endpoint must not be asked to launch again") +} + +// TestDynamicEndpointResolver_ReadyWithoutBaseURLDoesNotResolve guards against +// routing inference at an empty address: READY without a base URL is a record +// we cannot use, not one we can. +func TestDynamicEndpointResolver_ReadyWithoutBaseURLDoesNotResolve(t *testing.T) { + svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: StateReady}} + r := NewDynamicEndpointResolver(svc, nil, 0) + + baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, baseURL) +} + +// TestDynamicEndpointResolver_AbsentRequestsLaunch is the whole point of the +// change: a cold model asks the daemon for a VM and reports "not yet", which +// the invoke paths turn into a retryable ModelNotReadyException. +func TestDynamicEndpointResolver_AbsentRequestsLaunch(t *testing.T) { + svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: StateAbsent}} + r := NewDynamicEndpointResolver(svc, nil, 0) + + _, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.NoError(t, err) + assert.False(t, ok) + assert.Equal(t, int64(1), svc.ensureCalls.Load()) +} + +// TestDynamicEndpointResolver_InFlightStatesDoNotReEnsure covers the states +// where something is already happening: asking again changes nothing, and a +// draining endpoint must not be resurrected mid-teardown. +func TestDynamicEndpointResolver_InFlightStatesDoNotReEnsure(t *testing.T) { + for _, state := range []EndpointState{StateStarting, StateDraining} { + svc := &fakeEndpointService{record: EndpointRecord{ModelID: resolverTestModel, State: state}} + r := NewDynamicEndpointResolver(svc, nil, 0) + + _, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.NoError(t, err, "state %s", state) + assert.False(t, ok, "state %s", state) + assert.Zero(t, svc.ensureCalls.Load(), "state %s must not request a launch", state) + } +} + +// TestDynamicEndpointResolver_DescribeErrorPropagates keeps a broken control +// plane distinguishable from a cold model: an unreachable daemon is an error, +// not an indefinite "not ready". +func TestDynamicEndpointResolver_DescribeErrorPropagates(t *testing.T) { + svc := &fakeEndpointService{describeErr: errors.New("no responders")} + r := NewDynamicEndpointResolver(svc, nil, 0) + + _, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.Error(t, err) + assert.False(t, ok) + assert.Zero(t, svc.ensureCalls.Load()) +} + +// TestDynamicEndpointResolver_EnsureErrorPropagates preserves the daemon's own +// refusal, which carries the real AWS code: no GPU it can admit comes back as +// ModelNotReadyException with the daemon's message rather than a bare one. +func TestDynamicEndpointResolver_EnsureErrorPropagates(t *testing.T) { + svc := &fakeEndpointService{ + record: EndpointRecord{ModelID: resolverTestModel, State: StateAbsent}, + ensureErr: errors.New(awserrors.ErrorModelNotReadyException), + } + r := NewDynamicEndpointResolver(svc, nil, 0) + + _, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) + assert.False(t, ok) +} + +// TestDynamicEndpointResolver_CachesReadyBaseURL keeps a warm model off the +// bus: the steady state is every invoke hitting this path. +func TestDynamicEndpointResolver_CachesReadyBaseURL(t *testing.T) { + svc := &fakeEndpointService{record: readyRecord("http://10.0.0.9:8000")} + r := NewDynamicEndpointResolver(svc, nil, time.Minute) + + for range 3 { + baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "http://10.0.0.9:8000", baseURL) + } + assert.Equal(t, int64(1), svc.describeCalls.Load(), "a cached endpoint must not re-describe") +} + +// TestDynamicEndpointResolver_CacheExpires bounds how stale a base URL can +// get, which is what makes a deleted endpoint recoverable without an +// invalidation hook the resolver interface has no room for. +func TestDynamicEndpointResolver_CacheExpires(t *testing.T) { + svc := &fakeEndpointService{record: readyRecord("http://10.0.0.9:8000")} + r := NewDynamicEndpointResolver(svc, nil, time.Millisecond) + + _, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.NoError(t, err) + require.True(t, ok) + + svc.setRecord(readyRecord("http://10.0.0.10:8000")) + time.Sleep(5 * time.Millisecond) + + baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "http://10.0.0.10:8000", baseURL, "an expired entry must re-resolve") + assert.Equal(t, int64(2), svc.describeCalls.Load()) +} + +// TestDynamicEndpointResolver_ConcurrentColdCallsCollapse covers the burst a +// cold model actually sees: many in-flight invokes, one describe and one +// ensure between them. +func TestDynamicEndpointResolver_ConcurrentColdCallsCollapse(t *testing.T) { + release := make(chan struct{}) + svc := &fakeEndpointService{ + record: EndpointRecord{ModelID: resolverTestModel, State: StateAbsent}, + beforeDescribe: func() { <-release }, + } + r := NewDynamicEndpointResolver(svc, nil, 0) + + var wg sync.WaitGroup + for range 8 { + wg.Go(func() { + _, ok, err := r.Endpoint(context.Background(), resolverTestModel) + assert.NoError(t, err) + assert.False(t, ok) + }) + } + + // Hold the first describe until every caller has had a chance to join it, + // so the assertion below is about single-flight and not about scheduling. + time.Sleep(20 * time.Millisecond) + close(release) + wg.Wait() + + assert.Equal(t, int64(1), svc.describeCalls.Load()) + assert.Equal(t, int64(1), svc.ensureCalls.Load()) +} diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index 85f278951..c1730de34 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -21,6 +21,7 @@ import ( gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" gateway_ecr "github.com/mulgadc/spinifex/spinifex/gateway/ecr" gateway_ecrauth "github.com/mulgadc/spinifex/spinifex/gateway/ecrauth" + handlers_bedrock "github.com/mulgadc/spinifex/spinifex/handlers/bedrock" "github.com/mulgadc/spinifex/spinifex/handlers/ecr" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" handlers_quota "github.com/mulgadc/spinifex/spinifex/handlers/quota" @@ -327,10 +328,12 @@ func launchService(config *config.ClusterConfig) error { "accountID", admin.DefaultAccountID(), "models", len(gateway_bedrock.CatalogModelIDs())) } - // Bedrock self-host endpoints are pinned, so their - // OpenAI-compatible base URLs come from static config. OCHRE_VLLM_ENDPOINTS - // is a comma-separated list of modelId=baseURL pairs. + // 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. bedrockEndpoints := parseBedrockEndpoints(os.Getenv("OCHRE_VLLM_ENDPOINTS")) + bedrockEndpointResolver := handlers_bedrock.NewDynamicEndpointResolver( + handlers_bedrock.NewNATSEndpointService(natsConn), bedrockEndpoints, 0) // Bedrock invocation records: every Converse/InvokeModel call (streaming // or not) is published to the invocation stream, then fanned out by @@ -362,29 +365,30 @@ func launchService(config *config.ClusterConfig) error { go gateway_bedrock.NewUsageConsumer(bedrockUsage, bedrockPrices).Run(janitorCtx, usageConsumer) gw := gateway.GatewayConfig{ - Debug: nodeConfig.AWSGW.Debug, - DisableLogging: false, - NATSConn: natsConn, - Config: nodeConfig.AWSGW.Config, - ExpectedNodes: len(config.Nodes), - Region: nodeConfig.Region, - InternalSuffix: config.AWS.InternalSuffix, - RegistryPort: registryPort, - RegistryHost: registryHost, - AZ: nodeConfig.AZ, - IAMService: iamService, - STSService: stsService, - Version: version, - Commit: commit, - ECRRegistry: ecrRegistry, - ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, ecrAudience), - ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, ecrAudience), - BedrockCredentials: bedrockCredentials, - BedrockEndpoints: bedrockEndpoints, - BedrockLoggingConfig: bedrockLoggingConfig, - BedrockRecorder: bedrockRecorder, - BedrockAccess: bedrockAccess, - BedrockAccessAdmin: bedrockAccess, + Debug: nodeConfig.AWSGW.Debug, + DisableLogging: false, + NATSConn: natsConn, + Config: nodeConfig.AWSGW.Config, + ExpectedNodes: len(config.Nodes), + Region: nodeConfig.Region, + InternalSuffix: config.AWS.InternalSuffix, + RegistryPort: registryPort, + RegistryHost: registryHost, + AZ: nodeConfig.AZ, + IAMService: iamService, + STSService: stsService, + Version: version, + Commit: commit, + ECRRegistry: ecrRegistry, + ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, ecrAudience), + ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, ecrAudience), + BedrockCredentials: bedrockCredentials, + BedrockEndpoints: bedrockEndpoints, + BedrockEndpointResolver: bedrockEndpointResolver, + BedrockLoggingConfig: bedrockLoggingConfig, + BedrockRecorder: bedrockRecorder, + BedrockAccess: bedrockAccess, + BedrockAccessAdmin: bedrockAccess, } // Rotate the ECR signing key on a 30-day cadence, retaining the previous keys