mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 17:11:30 +00:00
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:
@@ -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".
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user