fix: cover the embedder cold-start window with retry and Retry-After (#934)

* fix: cover the embedder cold-start window with retry and Retry-After

Guardrail and knowledge-base embed calls hard-failed with a raw 503
whenever the embedding endpoint's connection was refused mid-relaunch,
instead of riding out the short window until it accepted connections.

embedBatch now retries a bounded number of times on connection-refused
or connection-reset, and any resulting 503 carries a Retry-After header
so callers back off instead of hammering the endpoint. A background
warm-up probe gates the embed path for statically pinned endpoints,
logging once when the endpoint first becomes reachable. Guardrail and
embedder infra logs now carry service/action labels for easier triage.

Fail-closed guardrail semantics are unchanged: an unavailable embedder
still denies unverified content, this only changes how quickly it
recovers once the endpoint is reachable again.

* fix: re-verify a persisted READY endpoint's liveness before resolving it

DynamicEndpointResolver's StateReady branch trusted the KV record's base
URL unconditionally. After an awsgw restart the record can still say
READY while the co-resident engine on that VM is mid-rebind, so the
resolver handed callers an address nothing was listening on yet.

describeAndEnsure now re-probes a READY member's own readiness route
before trusting it, short-TTL cached so a warm model still resolves at
no extra cost per request. A member that is not answering falls through
to the same not-ready path as a launch already in flight, so guardrail
and embed callers get the existing retryable answer instead of dialing
a dead address. This applies at the resolver seam, so it covers every
co-served member, not just the embedder.
This commit is contained in:
Josh Siviour
2026-08-28 13:04:25 +10:00
committed by GitHub
parent 3f55b89cd1
commit fb499a8c6a
12 changed files with 773 additions and 30 deletions
+41
View File
@@ -5,6 +5,7 @@ import (
"fmt"
"net/http"
"strings"
"time"
)
type ErrorMessage struct {
@@ -635,6 +636,46 @@ func Errorf(code, format string, args ...any) error {
return outer
}
// retryableError wraps a registered code with a suggested Retry-After
// duration, for a 503 that a client can expect to clear on its own within a
// bounded window (e.g. a warm-up race) rather than an open-ended outage.
// Error() returns only the code, so code resolution is unaffected.
type retryableError struct {
code string
retryAfter time.Duration
}
func (e *retryableError) Error() string { return e.code }
// RetryAfter wraps code (a registered ErrorLookup key) with a suggested
// Retry-After duration for ErrorHandler to surface as a response header.
func RetryAfter(code string, d time.Duration) error {
return &retryableError{code: code, retryAfter: d}
}
// ResolveRetryAfter returns the Retry-After duration attached to err via
// RetryAfter, if any, walking the same unwrap tree resolveErrorDetail does.
func ResolveRetryAfter(err error) (time.Duration, bool) {
if err == nil {
return 0, false
}
if re, ok := errors.AsType[*retryableError](err); ok {
return re.retryAfter, true
}
if joined, isJoined := err.(interface{ Unwrap() []error }); isJoined {
for _, inner := range joined.Unwrap() {
if d, found := ResolveRetryAfter(inner); found {
return d, true
}
}
return 0, false
}
if wrapped, isWrapped := err.(interface{ Unwrap() error }); isWrapped {
return ResolveRetryAfter(wrapped.Unwrap())
}
return 0, false
}
// errorLookupByService overrides ErrorLookup's message and HTTP status for a
// (service, code) pair whose wire code is shared by services with different
// canonical wording — e.g. ACM and EKS both use "ResourceInUseException".
+28
View File
@@ -5,6 +5,7 @@ import (
"fmt"
"strings"
"testing"
"time"
)
// TestErrorLookup_Structure asserts ErrorLookup invariants: minimum size, valid HTTP codes, non-empty messages.
@@ -138,6 +139,33 @@ func TestResolveErrorDetail_BareCodeCarriesNoMessage(t *testing.T) {
}
}
// TestResolveRetryAfter_WrappedAndResolvable covers the embedder warm-up
// case: RetryAfter's duration survives both a %w wrapper and code
// resolution, so ErrorHandler can set the header and still map the code.
func TestResolveRetryAfter_WrappedAndResolvable(t *testing.T) {
err := fmt.Errorf("embed: %w", RetryAfter(ErrorServiceUnavailableException, 2*time.Second))
d, ok := ResolveRetryAfter(err)
if !ok || d != 2*time.Second {
t.Fatalf("ResolveRetryAfter() = (%v, %v), want (2s, true)", d, ok)
}
code, _, ok := ResolveErrorDetail(err)
if !ok || code != ErrorServiceUnavailableException {
t.Fatalf("ResolveErrorDetail() code = (%q, %v), want (%q, true)", code, ok, ErrorServiceUnavailableException)
}
}
// TestResolveRetryAfter_PlainErrorNotFound is the compatibility case: an
// ordinary errors.New(code) -- the overwhelming majority of call sites --
// carries no Retry-After hint.
func TestResolveRetryAfter_PlainErrorNotFound(t *testing.T) {
_, ok := ResolveRetryAfter(errors.New(ErrorServiceUnavailableException))
if ok {
t.Errorf("ResolveRetryAfter() ok = true for a plain error, want false")
}
}
// TestLookupErrorMessage covers per-service scoping: ACM must not inherit
// EKS's wording for the shared ResourceInUseException code, and an unscoped
// pair must fall back to ErrorLookup's existing global default unchanged.
+111
View File
@@ -0,0 +1,111 @@
package gateway_bedrock
import (
"context"
"log/slog"
"net/http"
"sync"
"time"
)
// embedderHealthPath is the readiness route embeddingsProvider's warm-up
// probe polls -- TEI has no other health route, mirroring
// handlers/bedrock/readiness.go's default (non-vLLM) family selection.
const embedderHealthPath = "/health"
// embedderWarmupPollInterval spaces the background warm-up probe's checks,
// and doubles as the positive-cache TTL: a ready result is trusted between
// probes, then re-verified next tick. Short enough that a cold-started TEI
// is picked up quickly, long enough not to hammer a healthy endpoint.
const embedderWarmupPollInterval = 2 * time.Second
// embedderWarmupProbeTimeout bounds a single warm-up health check, well
// under embedderWarmupPollInterval, so a slow or hung endpoint cannot stall
// the background loop's cadence.
const embedderWarmupProbeTimeout = 1 * time.Second
// warmupProbe tracks whether a configured static embedder endpoint is
// currently answering its health route, so Embed can fail fast with a clean
// "not ready" signal instead of dialing a port TEI hasn't bound yet. A
// background goroutine keeps the cached state fresh for the process
// lifetime and logs once when the endpoint first becomes ready, giving
// operators an "embedder ready" marker to bound the cold-start window from
// the awsgw side, independent of request traffic.
type warmupProbe struct {
client *http.Client
baseURL string
mu sync.RWMutex
ready bool
}
// newWarmupProbeWithInterval constructs a warmupProbe for baseURL and starts
// its background poll loop at pollInterval, then returns it. Production
// callers pass embedderWarmupPollInterval (see embeddingsProvider.warmupFor);
// tests pass a short interval to drive the loop without waiting out the
// production cadence.
func newWarmupProbeWithInterval(client *http.Client, baseURL string, pollInterval time.Duration) *warmupProbe {
p := &warmupProbe{client: client, baseURL: baseURL}
go p.run(pollInterval)
return p
}
// run polls baseURL's health route every pollInterval for the life of the
// process, updating the cached ready state and logging exactly once on the
// not-ready -> ready transition.
func (p *warmupProbe) run(pollInterval time.Duration) {
p.probe()
ticker := time.NewTicker(pollInterval)
defer ticker.Stop()
for range ticker.C {
p.probe()
}
}
// probe issues one health check and updates the cached state, logging an
// "embedder ready" marker the first time it observes success.
func (p *warmupProbe) probe() {
ok := probeEmbedderHealth(p.client, p.baseURL)
p.mu.Lock()
wasReady := p.ready
p.ready = ok
p.mu.Unlock()
if ok && !wasReady {
slog.Info("embedder ready", "service", embedServiceLabel, "action", "warmup_probe", "endpoint", p.baseURL)
}
}
// Ready reports the most recently probed readiness state. It never blocks
// on network I/O -- the background loop in run does that -- so a caller on
// the request path gets an instant, if up-to-embedderWarmupPollInterval-
// stale, answer.
func (p *warmupProbe) Ready() bool {
p.mu.RLock()
defer p.mu.RUnlock()
return p.ready
}
// probeEmbedderHealth issues a single bounded GET against baseURL+
// embedderHealthPath and reports whether it returned HTTP 200. Any
// transport error (connection refused while TEI is still starting) or
// non-200 status counts as not-yet-ready, mirroring
// handlers/bedrock/readiness.go's probeOnce -- duplicated rather than
// imported, since handlers_bedrock already imports gateway_bedrock and Go
// forbids the reverse.
func probeEmbedderHealth(client *http.Client, baseURL string) bool {
ctx, cancel := context.WithTimeout(context.Background(), embedderWarmupProbeTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+embedderHealthPath, nil)
if err != nil {
return false
}
resp, err := client.Do(req)
if err != nil {
return false
}
defer resp.Body.Close()
return resp.StatusCode == http.StatusOK
}
@@ -0,0 +1,107 @@
// Exercises warmupProbe's background readiness cache directly, with no
// exported surface to drive it through.
//
//test:in-package
package gateway_bedrock
import (
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
"time"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestWarmupProbe_NotReadyUntilHealthPasses(t *testing.T) {
var healthy atomic.Bool
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, embedderHealthPath, r.URL.Path)
if healthy.Load() {
w.WriteHeader(http.StatusOK)
return
}
w.WriteHeader(http.StatusServiceUnavailable)
}))
defer ts.Close()
p := newWarmupProbeWithInterval(ts.Client(), ts.URL, 10*time.Millisecond)
// The immediate first probe run issues before the ticker starts lands
// against the not-yet-healthy server; give it a moment to complete.
time.Sleep(30 * time.Millisecond)
assert.False(t, p.Ready(), "must report not-ready until /health passes")
healthy.Store(true)
require.Eventually(t, p.Ready, time.Second, 5*time.Millisecond,
"must become ready once /health starts returning 200")
}
func TestWarmupProbe_ReProbesAfterInterval(t *testing.T) {
var healthy atomic.Bool
healthy.Store(true)
var calls atomic.Int64
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls.Add(1)
if healthy.Load() {
w.WriteHeader(http.StatusOK)
return
}
w.WriteHeader(http.StatusServiceUnavailable)
}))
defer ts.Close()
p := newWarmupProbeWithInterval(ts.Client(), ts.URL, 10*time.Millisecond)
require.Eventually(t, p.Ready, time.Second, 5*time.Millisecond,
"must become ready once the endpoint is healthy")
callsAtReady := calls.Load()
healthy.Store(false)
require.Eventually(t, func() bool { return !p.Ready() }, time.Second, 5*time.Millisecond,
"must re-probe on the next tick and flip back to not-ready")
assert.Greater(t, calls.Load(), callsAtReady, "must have issued at least one more probe after the first ready result")
}
func TestEmbeddingsProvider_Embed_WarmupGateFailsFastUntilHealthy(t *testing.T) {
var healthy atomic.Bool
var embedCalls atomic.Int64
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case embedderHealthPath:
if healthy.Load() {
w.WriteHeader(http.StatusOK)
} else {
w.WriteHeader(http.StatusServiceUnavailable)
}
case embeddingsPath:
embedCalls.Add(1)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1]}],"model":"nomic-embed-text-v1.5","usage":{"prompt_tokens":1,"total_tokens":1}}`))
default:
w.WriteHeader(http.StatusNotFound)
}
}))
defer ts.Close()
modelID := DefaultEmbeddingModel
p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
p.httpClient = ts.Client()
p.warmupPollInterval = 10 * time.Millisecond
p.warmupFor(ts.URL) // mirrors NewEmbedder(resolver, ts.URL)
time.Sleep(30 * time.Millisecond)
_, err := p.Embed(t.Context(), modelID, []string{"hello"})
require.Error(t, err)
assert.Equal(t, awserrors.ErrorServiceUnavailableException, err.Error())
assert.Zero(t, embedCalls.Load(), "must not dial the embeddings endpoint while the warm-up gate reports not ready")
healthy.Store(true)
require.Eventually(t, func() bool {
_, err := p.Embed(t.Context(), modelID, []string{"hello"})
return err == nil
}, time.Second, 10*time.Millisecond, "must succeed once the warm-up probe catches up to the endpoint being healthy")
}
+149 -20
View File
@@ -9,6 +9,8 @@ import (
"log/slog"
"net/http"
"sync"
"syscall"
"time"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
@@ -26,6 +28,22 @@ const (
// maxIntrospectionResponseBytes bounds the /info and /tokenize response
// bodies read into memory, mirroring the ingest path's bounded reads.
maxIntrospectionResponseBytes = 1 << 20
// embedServiceLabel tags every slog call in this file, so bedrock traffic
// filtered by labels.service in the log sink surfaces the embedder's own
// failure cause instead of only the generic gateway request wrapper.
embedServiceLabel = "bedrock-embeddings"
// embedTransportRetries bounds embedBatch's connection-refused/reset
// retry to this many extra attempts after the first, spread over
// ~1.5s by embedTransportRetryBackoff -- enough to catch a request that
// lands just before TEI's listener binds, without turning one cold
// request into a multi-second stall.
embedTransportRetries = 3
// embedTransportRetryBackoff is the fixed pause between embedBatch's
// transport-error retries.
embedTransportRetryBackoff = 500 * time.Millisecond
)
// DefaultEmbeddingModel is Ochre's GPU-served embedding model, co-resident in
@@ -101,6 +119,20 @@ type embeddingsProvider struct {
// changes for a running server, so this is cached for the process
// lifetime rather than re-fetched per call.
infoCache sync.Map
// warmups holds baseURL -> *warmupProbe for the subset of resolved
// endpoints NewEmbedder was told to warm-up-gate (the static
// OCHRE_VLLM_ENDPOINTS bypass -- see NewEmbedder). A baseURL absent from
// this map is never gated: the daemon's dynamic endpoint registry
// already confirms readiness itself before handing back a base URL, so
// gating there would just add a redundant, possibly stale, delay.
warmupMu sync.Mutex
warmups map[string]*warmupProbe
// warmupPollInterval overrides embedderWarmupPollInterval for a probe
// created by warmupFor. Zero (the default) keeps the production
// interval; tests set this before the first warmupFor call to drive the
// background loop on a short cadence.
warmupPollInterval time.Duration
}
var _ Embedder = (*embeddingsProvider)(nil)
@@ -120,8 +152,53 @@ func newEmbeddingsProvider(endpointResolver EndpointResolver) *embeddingsProvide
// NewEmbedder is newEmbeddingsProvider's exported constructor, for callers
// outside this package (e.g. the daemon's Ochre vector store wiring) that
// need an Embedder without reaching into gateway_bedrock's unexported types.
func NewEmbedder(endpointResolver EndpointResolver) Embedder {
return newEmbeddingsProvider(endpointResolver)
// warmupBaseURLs lists resolved base URLs that bypass the daemon's own
// readiness lifecycle -- in practice the OCHRE_VLLM_ENDPOINTS static pin for
// the embedding model -- and so need their own background /health warm-up
// probe before Embed dials them. Omit it (the zero-arg call) to keep today's
// behaviour: no gating at all.
func NewEmbedder(endpointResolver EndpointResolver, warmupBaseURLs ...string) Embedder {
p := newEmbeddingsProvider(endpointResolver)
for _, baseURL := range warmupBaseURLs {
if baseURL == "" {
continue
}
p.warmupFor(baseURL)
}
return p
}
// warmupFor returns (creating and starting if needed) the background
// readiness prober for baseURL, so a cold endpoint's health is tracked once
// regardless of how many callers register it.
func (p *embeddingsProvider) warmupFor(baseURL string) *warmupProbe {
p.warmupMu.Lock()
defer p.warmupMu.Unlock()
if p.warmups == nil {
p.warmups = make(map[string]*warmupProbe)
}
if wp, ok := p.warmups[baseURL]; ok {
return wp
}
interval := p.warmupPollInterval
if interval <= 0 {
interval = embedderWarmupPollInterval
}
wp := newWarmupProbeWithInterval(p.httpClient, baseURL, interval)
p.warmups[baseURL] = wp
return wp
}
// trackedWarmup returns baseURL's registered warm-up prober, if any. false
// means baseURL is not gated -- either it never came through
// NewEmbedder's warmupBaseURLs, or this embeddingsProvider was built via
// newEmbeddingsProvider directly (tests, and any caller that wants today's
// ungated behaviour).
func (p *embeddingsProvider) trackedWarmup(baseURL string) (*warmupProbe, bool) {
p.warmupMu.Lock()
defer p.warmupMu.Unlock()
wp, ok := p.warmups[baseURL]
return wp, ok
}
// maxEmbedClientBatch bounds inputs per embeddings request to TEI's default
@@ -140,11 +217,23 @@ func (p *embeddingsProvider) Embed(ctx context.Context, modelID string, inputs [
baseURL, ok, err := p.endpointResolver.Endpoint(ctx, modelID)
if err != nil {
slog.Error("embeddings: endpoint resolution failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
slog.Error("embeddings: endpoint resolution failed", "service", embedServiceLabel, "action", "resolve_endpoint", "model", modelID, "err", err)
return nil, awserrors.RetryAfter(awserrors.ErrorServiceUnavailableException, embedderWarmupPollInterval)
}
if !ok {
return nil, errors.New(awserrors.ErrorModelNotReadyException)
// Includes the resolver's own liveness gate on a stale READY record:
// same retryable signal as a model that has not launched at all.
return nil, awserrors.RetryAfter(awserrors.ErrorModelNotReadyException, embedderWarmupPollInterval)
}
// A resolved baseURL under warm-up gating (the static endpoint bypass --
// see NewEmbedder) that hasn't yet answered its health probe fails fast
// here instead of dialing a port TEI hasn't bound yet. embedBatch's own
// transport retry (below) still covers the narrow race right as the
// background probe catches up.
if wp, tracked := p.trackedWarmup(baseURL); tracked && !wp.Ready() {
slog.Warn("embeddings: endpoint not ready, failing closed during warm-up", "service", embedServiceLabel, "action", "embed", "model", modelID, "endpoint", baseURL)
return nil, awserrors.RetryAfter(awserrors.ErrorServiceUnavailableException, embedderWarmupPollInterval)
}
out := make([][]float32, 0, len(inputs))
@@ -165,49 +254,89 @@ func (p *embeddingsProvider) Embed(ctx context.Context, modelID string, inputs [
func (p *embeddingsProvider) embedBatch(ctx context.Context, modelID, baseURL string, inputs []string) ([][]float32, error) {
reqBody, err := json.Marshal(embeddingsRequest{Model: modelID, Input: inputs})
if err != nil {
slog.Error("embeddings: failed to marshal request", "model", modelID, "err", err)
slog.Error("embeddings: failed to marshal request", "service", embedServiceLabel, "action", "embed_batch", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+embeddingsPath, bytes.NewReader(reqBody))
resp, err := p.postEmbedRequest(ctx, modelID, baseURL, reqBody)
if err != nil {
slog.Error("embeddings: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq.Header.Set("Content-Type", "application/json")
resp, err := p.httpClient.Do(httpReq)
if err != nil {
slog.Error("embeddings: request failed", "model", modelID, "endpoint", baseURL, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
return nil, err
}
defer resp.Body.Close()
respBody, err := io.ReadAll(resp.Body)
if err != nil {
slog.Error("embeddings: failed to read response body", "model", modelID, "err", err)
slog.Error("embeddings: failed to read response body", "service", embedServiceLabel, "action", "embed_batch", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
slog.Error("embeddings: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
slog.Error("embeddings: upstream error", "service", embedServiceLabel, "action", "embed_batch", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
var er embeddingsResponse
if err := json.Unmarshal(respBody, &er); err != nil {
slog.Error("embeddings: failed to parse response", "model", modelID, "err", err)
slog.Error("embeddings: failed to parse response", "service", embedServiceLabel, "action", "embed_batch", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorModelErrorException)
}
vectors, err := orderEmbeddings(len(inputs), er.Data)
if err != nil {
slog.Error("embeddings: malformed response shape", "model", modelID, "err", err)
slog.Error("embeddings: malformed response shape", "service", embedServiceLabel, "action", "embed_batch", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorModelErrorException)
}
return vectors, nil
}
// postEmbedRequest POSTs reqBody to baseURL+embeddingsPath, retrying a
// connection-refused or connection-reset transport error up to
// embedTransportRetries times with a fixed backoff -- the shape of a
// request landing in the window between an awsgw restart and TEI's listener
// binding. Any other failure (a reached but erroring endpoint, a timeout,
// ctx cancellation) is not retried: retrying those would only add latency
// to a call that was never going to succeed.
func (p *embeddingsProvider) postEmbedRequest(ctx context.Context, modelID, baseURL string, reqBody []byte) (*http.Response, error) {
var lastErr error
for attempt := 0; attempt <= embedTransportRetries; attempt++ {
if attempt > 0 {
select {
case <-ctx.Done():
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
case <-time.After(embedTransportRetryBackoff):
}
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+embeddingsPath, bytes.NewReader(reqBody))
if err != nil {
slog.Error("embeddings: failed to build request", "service", embedServiceLabel, "action", "embed_batch", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq.Header.Set("Content-Type", "application/json")
resp, doErr := p.httpClient.Do(httpReq)
if doErr == nil {
return resp, nil
}
lastErr = doErr
if !isRetryableTransportError(doErr) || attempt == embedTransportRetries {
break
}
slog.Warn("embeddings: transport error, retrying", "service", embedServiceLabel, "action", "embed_batch", "model", modelID, "endpoint", baseURL, "attempt", attempt+1, "err", doErr)
}
slog.Error("embeddings: request failed", "service", embedServiceLabel, "action", "embed_batch", "model", modelID, "endpoint", baseURL, "err", lastErr)
return nil, awserrors.RetryAfter(awserrors.ErrorServiceUnavailableException, embedderWarmupPollInterval)
}
// isRetryableTransportError reports whether err is a connection-refused or
// connection-reset failure -- dialing a TEI port that hasn't bound yet, or
// one torn down mid-request -- as opposed to a reached-but-erroring
// endpoint, a timeout, or a malformed request, none of which retrying helps.
func isRetryableTransportError(err error) bool {
return errors.Is(err, syscall.ECONNREFUSED) || errors.Is(err, syscall.ECONNRESET)
}
// MaxInputLength implements TokenLimiter: it resolves modelID's endpoint and
// returns its cached (or freshly fetched) TEI /info max_input_length,
// falling back to defaultMaxInputLength whenever the endpoint can't be
@@ -7,9 +7,12 @@ package gateway_bedrock
import (
"context"
"encoding/json"
"errors"
"fmt"
"net"
"net/http"
"net/http/httptest"
"syscall"
"testing"
"github.com/mulgadc/spinifex/spinifex/awserrors"
@@ -354,3 +357,81 @@ func TestEmbeddingsProvider_CountTokens_ResolverMissReturnsNotOK(t *testing.T) {
assert.False(t, ok)
assert.Zero(t, count)
}
// countingErrTransport fails the first failures RoundTrip calls with err,
// then delegates to next -- lets a test control exactly how many transport
// failures embedBatch's retry must absorb before (or whether) it succeeds.
type countingErrTransport struct {
failures int
err error
next http.RoundTripper
calls int
}
func (t *countingErrTransport) RoundTrip(req *http.Request) (*http.Response, error) {
t.calls++
if t.calls <= t.failures {
return nil, t.err
}
return t.next.RoundTrip(req)
}
func TestEmbeddingsProvider_EmbedBatch_RetriesTransportErrorThenSucceeds(t *testing.T) {
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(`{"data":[{"index":0,"embedding":[0.1]}],"model":"nomic-embed-text-v1.5","usage":{"prompt_tokens":1,"total_tokens":1}}`))
}))
defer ts.Close()
transport := &countingErrTransport{
failures: embedTransportRetries, // fail every attempt but the last
err: &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED},
next: ts.Client().Transport,
}
modelID := DefaultEmbeddingModel
p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
p.httpClient = &http.Client{Transport: transport}
vectors, err := p.Embed(context.Background(), modelID, []string{"hello"})
require.NoError(t, err)
require.Len(t, vectors, 1)
assert.Equal(t, embedTransportRetries+1, transport.calls,
"must retry the connection-refused transport error up to the cap before succeeding")
}
func TestEmbeddingsProvider_EmbedBatch_GivesUpAfterTransportRetryCap(t *testing.T) {
transport := &countingErrTransport{
failures: embedTransportRetries + 10, // always fail
err: &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED},
next: http.DefaultTransport,
}
modelID := DefaultEmbeddingModel
p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: "http://127.0.0.1:1"}))
p.httpClient = &http.Client{Transport: transport}
_, err := p.Embed(context.Background(), modelID, []string{"hello"})
require.Error(t, err)
assert.Equal(t, awserrors.ErrorServiceUnavailableException, err.Error())
assert.Equal(t, embedTransportRetries+1, transport.calls,
"must give up after embedTransportRetries+1 attempts, not retry forever")
}
func TestEmbeddingsProvider_EmbedBatch_NonTransportErrorNotRetried(t *testing.T) {
transport := &countingErrTransport{
failures: embedTransportRetries + 10, // always fail
err: errors.New("boom: not a transport error"),
next: http.DefaultTransport,
}
modelID := DefaultEmbeddingModel
p := newEmbeddingsProvider(NewStaticEndpointResolver(map[string]string{modelID: "http://127.0.0.1:1"}))
p.httpClient = &http.Client{Transport: transport}
_, err := p.Embed(context.Background(), modelID, []string{"hello"})
require.Error(t, err)
assert.Equal(t, awserrors.ErrorServiceUnavailableException, err.Error())
assert.Equal(t, 1, transport.calls, "a non-transport error must not be retried")
}
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"fmt"
"log/slog"
"regexp"
"slices"
@@ -317,6 +318,11 @@ func assessTopicPolicy(ctx context.Context, embedder Embedder, cfg *bedrock.Guar
case blocked:
return assessment, true, nil
case unverified:
// Logged here rather than only at the embedder call sites, so every
// "topic policy failed closed" outcome carries the same operator
// signal regardless of which policy check ends up surfacing it.
slog.Warn("guardrail: topic policy unverified, failing closed",
"service", guardrailServiceLabel, "action", "topic_policy")
return assessment, false, errors.New(awserrors.ErrorServiceUnavailableException)
default:
return assessment, false, nil
@@ -19,6 +19,12 @@ import (
// here if live traffic shows the false positive/negative rate drifting.
const defaultTopicSimilarityThreshold = 0.42
// guardrailServiceLabel tags every slog call in this file, so guardrail
// traffic filtered by labels.service in the log sink surfaces the semantic
// (embedder-backed) topic check's own failure cause instead of only the
// generic gateway request wrapper.
const guardrailServiceLabel = "bedrock-guardrail"
// topicSimilarityThreshold returns topic's similarity threshold. AWS's
// GuardrailTopicConfig wire shape has no field to override this per topic,
// so every topic currently gets the same default; this indirection is the
@@ -92,7 +98,7 @@ func topicVectors(ctx context.Context, embedder Embedder, modelID string, topic
vectors, err := embedder.Embed(ctx, modelID, phrases)
if err != nil {
slog.Warn("guardrail: topic embedding unavailable, failing closed on unverified content",
"topic", aws.StringValue(topic.Name), "model", modelID, "err", err)
"service", guardrailServiceLabel, "action", "topic_vectors", "topic", aws.StringValue(topic.Name), "model", modelID, "err", err)
return nil, err
}
topicVectorCache.Store(key, vectors)
@@ -138,12 +144,14 @@ func maxCosineSimilarity(vector []float32, candidates [][]float32) float64 {
// caller must fail closed rather than silently pass unverified content.
func embedGuardrailTexts(ctx context.Context, embedder Embedder, modelID string, texts []string) ([][]float32, error) {
if embedder == nil {
slog.Warn("guardrail: no embedder configured, falling back to literal topic matching")
slog.Warn("guardrail: no embedder configured, falling back to literal topic matching",
"service", guardrailServiceLabel, "action", "embed_texts")
return nil, nil
}
vectors, err := embedder.Embed(ctx, modelID, texts)
if err != nil {
slog.Warn("guardrail: input embedding unavailable, failing closed on unverified content", "model", modelID, "err", err)
slog.Warn("guardrail: input embedding unavailable, failing closed on unverified content",
"service", guardrailServiceLabel, "action", "embed_texts", "model", modelID, "err", err)
return nil, err
}
return vectors, nil
+8
View File
@@ -764,6 +764,14 @@ func (gw *GatewayConfig) ErrorHandler(w http.ResponseWriter, r *http.Request, er
errorMsg.HTTPCode = 500
}
// A 503 the producing call site knows is a bounded, self-clearing
// condition (e.g. the embedder warm-up window) carries a suggested
// Retry-After so SDK clients back off on their own cadence instead of
// hammering the endpoint every request.
if d, ok := awserrors.ResolveRetryAfter(err); ok {
w.Header().Set("Retry-After", strconv.Itoa(int(d.Seconds())))
}
// EKS, ECR, ACM, ECS, tagging, and the bedrock family use AWS JSON 1.1;
// query/XML services fall through.
if jsonErrorService(svc) {
+94 -6
View File
@@ -2,6 +2,7 @@ package handlers_bedrock
import (
"context"
"net/http"
"sync"
"time"
@@ -21,6 +22,18 @@ const defaultEndpointCacheTTL = 5 * time.Second
// is to return the moment the endpoint answers.
const coldStartPollInterval = 500 * time.Millisecond
// defaultLivenessCacheTTL bounds how long a liveness re-probe of a persisted
// READY member is reused. Short, so a member that goes dark between probes
// (e.g. a co-resident engine still rebinding its port right after an awsgw
// restart) is caught within a couple of requests, but long enough above zero
// that resolving a warm, steady-state model costs one HTTP round trip per
// TTL rather than one per request.
const defaultLivenessCacheTTL = 2 * time.Second
// livenessProbeTimeout bounds a single liveness re-probe, so a member that
// accepts the connection but never answers cannot stall resolution.
const livenessProbeTimeout = 500 * time.Millisecond
// DynamicEndpointResolver resolves a self-host model's base URL through the
// daemon's endpoint registry, and asks for a launch when there is nothing to
// resolve. It lives here rather than in gateway_bedrock because that package
@@ -46,6 +59,16 @@ type DynamicEndpointResolver struct {
mu sync.Mutex
cached map[string]cachedEndpoint
// client issues liveness re-probes against a persisted READY member's own
// health route (see isLive). Not used for anything else here.
client *http.Client
// livenessTTL bounds how long a liveness probe result is reused; see
// defaultLivenessCacheTTL.
livenessTTL time.Duration
livenessMu sync.Mutex
livenessCache map[string]livenessResult
}
// cachedEndpoint is one resolved READY base URL and the moment it expires.
@@ -54,6 +77,13 @@ type cachedEndpoint struct {
expiresAt time.Time
}
// livenessResult is the outcome of the last liveness re-probe of one base
// URL and when it was taken.
type livenessResult struct {
live bool
checkedAt time.Time
}
var _ gateway_bedrock.EndpointResolver = (*DynamicEndpointResolver)(nil)
// ResolverOption adjusts a DynamicEndpointResolver at construction.
@@ -67,6 +97,19 @@ func WithColdStartWait(d time.Duration) ResolverOption {
return func(r *DynamicEndpointResolver) { r.coldStartWait = d }
}
// WithHTTPClient overrides the client isLive re-probes a READY member on.
// Production takes the default; tests point it at an httptest server.
func WithHTTPClient(c *http.Client) ResolverOption {
return func(r *DynamicEndpointResolver) { r.client = c }
}
// WithLivenessCacheTTL overrides how long a liveness probe result is reused,
// so a test can exercise re-verification without waiting out the production
// interval.
func WithLivenessCacheTTL(d time.Duration) ResolverOption {
return func(r *DynamicEndpointResolver) { r.livenessTTL = d }
}
// NewDynamicEndpointResolver builds a resolver over svc. Entries in static
// (OCHRE_VLLM_ENDPOINTS) are resolved first and never reach svc, so a pinned
// endpoint bypasses the lifecycle entirely. A zero ttl takes the default.
@@ -76,10 +119,13 @@ func NewDynamicEndpointResolver(svc EndpointService, static map[string]string, t
ttl = defaultEndpointCacheTTL
}
r := &DynamicEndpointResolver{
svc: svc,
static: static,
ttl: ttl,
cached: make(map[string]cachedEndpoint),
svc: svc,
static: static,
ttl: ttl,
cached: make(map[string]cachedEndpoint),
client: &http.Client{},
livenessTTL: defaultLivenessCacheTTL,
livenessCache: make(map[string]livenessResult),
}
for _, opt := range opts {
opt(r)
@@ -152,7 +198,15 @@ func (r *DynamicEndpointResolver) describeAndEnsure(ctx context.Context, account
switch out.Endpoint.State {
case StateReady:
if url := out.Endpoint.MemberBaseURL(modelID); url != "" {
return url, nil
if r.isLive(ctx, out.Endpoint, modelID, url) {
return url, nil
}
// The persisted record is READY but the member itself is not
// answering (e.g. a co-resident engine still rebinding its port
// right after an awsgw restart) -- nothing new to ensure, so
// this is the same "wait or report not-ready" shape as a launch
// already in flight, not a fresh launch request.
return r.awaitReady(ctx, accountID, modelID)
}
return "", nil
case StateStarting:
@@ -199,7 +253,7 @@ func (r *DynamicEndpointResolver) awaitReady(ctx context.Context, accountID, mod
return "", err
}
if out.Endpoint.State == StateReady {
if url := out.Endpoint.MemberBaseURL(modelID); url != "" {
if url := out.Endpoint.MemberBaseURL(modelID); url != "" && r.isLive(ctx, out.Endpoint, modelID, url) {
return url, nil
}
}
@@ -209,6 +263,40 @@ func (r *DynamicEndpointResolver) awaitReady(ctx context.Context, accountID, mod
}
}
// isLive re-probes a persisted READY member's own readiness route before the
// resolver trusts it, short-TTL cached per base URL so a warm, steady-state
// model costs one HTTP round trip per livenessTTL rather than per resolve.
// A record with no per-member entry predates the Members field (or is a test
// fixture) and has no route to probe on, so it is trusted as before.
//
// This exists because a persisted READY state only means the endpoint was
// observed serving at some point in the past: after an awsgw restart the KV
// record survives untouched while a co-resident engine on the same VM is
// still rebinding its port, and describeAndEnsure has no other signal that
// the address it is about to hand out currently dials to nothing.
func (r *DynamicEndpointResolver) isLive(ctx context.Context, rec EndpointRecord, modelID, baseURL string) bool {
member, ok := rec.Members[modelID]
if !ok {
return true
}
r.livenessMu.Lock()
if cached, ok := r.livenessCache[baseURL]; ok && time.Since(cached.checkedAt) < r.livenessTTL {
r.livenessMu.Unlock()
return cached.live
}
r.livenessMu.Unlock()
probeCtx, cancel := context.WithTimeout(ctx, livenessProbeTimeout)
defer cancel()
live := probeOnce(probeCtx, r.client, baseURL+readinessPath(member.Family))
r.livenessMu.Lock()
r.livenessCache[baseURL] = livenessResult{live: live, checkedAt: time.Now()}
r.livenessMu.Unlock()
return live
}
func (r *DynamicEndpointResolver) lookupCache(modelID string) (string, bool) {
r.mu.Lock()
defer r.mu.Unlock()
+124
View File
@@ -3,6 +3,10 @@ package handlers_bedrock
import (
"context"
"errors"
"net"
"net/http"
"net/http/httptest"
"net/url"
"sync"
"sync/atomic"
"testing"
@@ -431,3 +435,123 @@ func TestDynamicEndpointResolver_EndpointForAccount_DoesNotShareTheGlobalCache(t
assert.Equal(t, int64(2), svc.describeCalls.Load(),
"the account-scoped resolve must not have pre-populated Endpoint's cache")
}
// memberReadyRecord builds a StateReady record with a real Members entry,
// pointing at srv's own host:port, so isLive has a route to actually probe --
// unlike readyRecord's bare BaseURL, which isLive trusts without a probe.
func memberReadyRecord(t *testing.T, srv *httptest.Server) EndpointRecord {
t.Helper()
u, err := url.Parse(srv.URL)
require.NoError(t, err)
host, portStr, err := net.SplitHostPort(u.Host)
require.NoError(t, err)
port, err := net.LookupPort("tcp", portStr)
require.NoError(t, err)
return EndpointRecord{
ModelID: resolverTestModel,
State: StateReady,
PrivateIP: host,
Members: map[string]MemberEndpoint{resolverTestModel: {Port: port}},
}
}
// TestDynamicEndpointResolver_ReadyRecordNotLiveResolvesAsNotReady is the
// heart of the fix: a persisted READY record whose member does not actually
// answer its readiness route must not be handed out as a live address --
// the caller gets the same retryable not-ready answer as a cold model.
func TestDynamicEndpointResolver_ReadyRecordNotLiveResolvesAsNotReady(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.WriteHeader(http.StatusServiceUnavailable)
}))
defer srv.Close()
svc := &fakeEndpointService{record: memberReadyRecord(t, srv)}
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_ReadyRecordRecoversOnceLive covers the
// recovery half: once the member starts answering, a fresh resolve (past the
// liveness cache TTL) succeeds without the daemon ever touching the record.
func TestDynamicEndpointResolver_ReadyRecordRecoversOnceLive(t *testing.T) {
var live atomic.Bool
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
if live.Load() {
w.WriteHeader(http.StatusOK)
return
}
w.WriteHeader(http.StatusServiceUnavailable)
}))
defer srv.Close()
svc := &fakeEndpointService{record: memberReadyRecord(t, srv)}
r := NewDynamicEndpointResolver(svc, nil, 0, WithLivenessCacheTTL(10*time.Millisecond))
_, ok, err := r.Endpoint(context.Background(), resolverTestModel)
require.NoError(t, err)
assert.False(t, ok, "not yet live")
live.Store(true)
require.Eventually(t, func() bool {
baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel)
return err == nil && ok && baseURL == srv.URL
}, time.Second, 5*time.Millisecond, "must resolve once the member answers and the liveness cache expires")
}
// TestDynamicEndpointResolver_ReadyRecordLiveResolvesImmediately guards the
// happy path the liveness gate must not regress: a member that is actually
// answering resolves on the first call, no waiting.
func TestDynamicEndpointResolver_ReadyRecordLiveResolvesImmediately(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
svc := &fakeEndpointService{record: memberReadyRecord(t, srv)}
r := NewDynamicEndpointResolver(svc, nil, 0)
baseURL, ok, err := r.Endpoint(context.Background(), resolverTestModel)
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, srv.URL, baseURL)
}
// TestDynamicEndpointResolver_IsLiveCachesWithinTTLThenRechecks pins down the
// cache half directly: a probe result is reused within livenessTTL (the hot
// path for a warm, steady-state model costs no HTTP round trip per resolve),
// and a fresh probe runs once that TTL has elapsed.
func TestDynamicEndpointResolver_IsLiveCachesWithinTTLThenRechecks(t *testing.T) {
var probes atomic.Int64
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
probes.Add(1)
w.WriteHeader(http.StatusOK)
}))
defer srv.Close()
r := NewDynamicEndpointResolver(&fakeEndpointService{}, nil, 0, WithLivenessCacheTTL(50*time.Millisecond))
rec := memberReadyRecord(t, srv)
for range 3 {
assert.True(t, r.isLive(context.Background(), rec, resolverTestModel, srv.URL))
}
assert.Equal(t, int64(1), probes.Load(), "a cached liveness result must not re-probe")
time.Sleep(60 * time.Millisecond)
assert.True(t, r.isLive(context.Background(), rec, resolverTestModel, srv.URL))
assert.Equal(t, int64(2), probes.Load(), "an expired cache entry must re-probe")
}
// TestDynamicEndpointResolver_IsLiveTrustsRecordsWithoutMembers keeps every
// pre-existing record shape (readyRecord, and any record predating the
// Members field) working exactly as before: with nothing to probe on, isLive
// trusts the state rather than manufacturing a route to check.
func TestDynamicEndpointResolver_IsLiveTrustsRecordsWithoutMembers(t *testing.T) {
r := NewDynamicEndpointResolver(&fakeEndpointService{}, nil, 0)
rec := readyRecord("http://10.0.0.9:8000")
assert.True(t, r.isLive(context.Background(), rec, resolverTestModel, rec.BaseURL))
}
+13 -1
View File
@@ -405,7 +405,19 @@ func launchService(config *config.ClusterConfig) error {
// Guardrail topicPolicy's semantic match reuses this same endpoint
// resolver (the one every self-hosted model, including the embedding
// model, already resolves through) rather than standing up a second one.
bedrockEmbedder := gateway_bedrock.NewEmbedder(bedrockEndpointResolver)
//
// A pinned OCHRE_VLLM_ENDPOINTS entry for the embedding model bypasses
// the daemon's own readiness lifecycle entirely (DynamicEndpointResolver
// resolves it from the static map before ever asking svc), so on
// restart awsgw has no signal that the co-resident TEI hasn't bound its
// port yet. Pass that one base URL through as a warm-up target so
// NewEmbedder background-probes it and fails closed cleanly instead of
// dialing a connection-refused port.
var bedrockEmbedderWarmupEndpoints []string
if baseURL, ok := bedrockEndpoints[gateway_bedrock.DefaultEmbeddingModel]; ok {
bedrockEmbedderWarmupEndpoints = append(bedrockEmbedderWarmupEndpoints, baseURL)
}
bedrockEmbedder := gateway_bedrock.NewEmbedder(bedrockEndpointResolver, bedrockEmbedderWarmupEndpoints...)
// Bedrock provisioned throughput: commitment metadata lives in the
// bedrock-provisioned KV bucket (gateway control plane), while the pinned