From 261d5d505e9efa07109bbe9495c19db7d057bebe Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Mon, 3 Aug 2026 12:34:07 +1000 Subject: [PATCH] feat: Ochre catalog serving spec and a servable self-host model (#672) * feat: in-tree serving specs + KV weights override for ochre catalog Extend catalogEntry with MinVRAMMiB, InstanceType and VLLMArgs so a self-host model's serving requirements ship with the catalog rather than living only in an operator's head. Add WeightsStore, a deployment-scoped (not account-scoped) JetStream KV override in a new bedrock-weights bucket that resolves a model's staged weights snapshot ID, following the same shape as the existing CredentialStore. It is plumbed into tieredCatalog and GetFoundationModel through a package-level resolver (SetWeightsResolver/currentWeightsResolver), since both are called from gateway/bedrock.go's fixed-arity route table, which is out of scope here. A self-host model with no resolvable weights snapshot is not servable and must not be advertised: tieredCatalog filters it out and GetFoundationModel returns ResourceNotFoundException for it. Replace the meta.llama3-70b-instruct-v1:0 self-host entry, which needs roughly 140 GB of VRAM in fp16 and has never actually run, with meta.llama3-2-1b-instruct-v1:0, a real AWS-published Bedrock model ID. At bf16 its 1.24B params are roughly 2.5 GiB of weights, well inside wattle's single RTX A1000 (8188 MiB); MinVRAMMiB is set to 4096 to budget headroom for CUDA context, activations and KV cache on top of that. This is an estimate, not a measurement off wattle - .7.5's launcher should correct it against a real nvidia-smi reading once the model is actually running there. lookupCatalogEntry, used by the untiered invoke/converse routing path, is left unchanged and still ignores weights resolvability; gating that path is out of scope for this bead. * fix: separate weights resolve errors from not-found in ochre catalog An internal fault (e.g. a JetStream outage) reaching the KV weights resolver was being collapsed into the same "not servable" verdict as a clean not-found: tieredCatalog would silently unadvertise every self-host model, and GetFoundationModel would return ResourceNotFoundException for a model that does exist but whose resolver is temporarily broken. Both call sites now distinguish the two: a resolve error is logged and surfaced as its own error (mapped to ServerInternal at the AWS error boundary), and only a clean not-resolvable result yields ResourceNotFoundException. tieredCatalog's signature changes to return an error alongside the entries so ListFoundationModels can propagate a resolve fault instead of silently returning a thinned-out list. Also raises the self-host entry's gpu-memory-utilization from 0.5 to 0.6 (roughly 4913 MiB of vLLM's own pool) and MinVRAMMiB from 4096 to 5120 to match: at 0.5, weights (~2.5 GiB) left under 1.5 GiB for KV cache and activations at max-model-len=8192, likely too tight for any real concurrency. The two numbers are meant to agree: MinVRAMMiB is the admission-gate floor, gpu-memory-utilization is what tells vLLM to actually claim that much on an 8188 MiB device. Comments trimmed to the repo's 3-line limit and stripped of plan/bead references; the VRAM reasoning and model-choice rationale live in the commit body and the plan doc instead. Threading the weights resolver through as an explicit parameter (mirroring CredentialResolver) rather than a package-level global was considered per review. Doing so needs edits to gateway/bedrock.go (bedrockRouteHandler's signature, both route closures, the Bedrock_Request call site, a new bedrockWeightsResolver() method) and gateway/gateway.go (a new BedrockWeights field on GatewayConfig). Both are outside gateway/bedrock/ and are explicitly out of scope for this change, owned by a parallel workstream already modifying the same route table. The global resolver stays, now with the error-handling fix above. * fix: stop e2e GetFoundationModel regressing on the new weights gate TestGetFoundationModel's "known model" subtest asserted on meta.llama3-2-1b-instruct-v1:0 unconditionally, but that entry is now gated on a weights snapshot staged in the bedrock-weights KV bucket. On any real deployment that bucket starts empty, so the subtest would fail with ResourceNotFoundException outside of OCHRE_E2E_SELFHOST. The "known model" assertion now targets the Anthropic entry, which resolves unconditionally and keeps testing what the subtest was actually there for: a known model round-trips, an unknown one is ResourceNotFoundException. The self-host assertion moves to a new TestGetFoundationModelSelfHost, gated behind OCHRE_E2E_SELFHOST=1 alongside TestConverseSelfHost, where weights are expected staged. tests/integration/bedrock_test.go's TestListFoundationModels already wires a weights-resolver stub before asserting the self-host entry is present, so it does not regress the same way. * test: correct a stale assertion message on the ochre catalog weights gate --- spinifex/gateway/bedrock/catalog.go | 64 ++++++-- spinifex/gateway/bedrock/catalog_test.go | 82 ++++++++++- spinifex/gateway/bedrock/contract_test.go | 4 +- .../gateway/bedrock/converse_stream_test.go | 6 +- spinifex/gateway/bedrock/invoke_llama_test.go | 14 +- .../gateway/bedrock/invoke_stream_test.go | 4 +- spinifex/gateway/bedrock/invoke_test.go | 10 +- spinifex/gateway/bedrock/router_test.go | 10 +- spinifex/gateway/bedrock/vllm_test.go | 16 +- spinifex/gateway/bedrock/weights.go | 138 ++++++++++++++++++ spinifex/gateway/bedrock/weights_test.go | 77 ++++++++++ spinifex/gateway/bedrock_request_test.go | 20 ++- spinifex/services/awsgw/awsgw.go | 7 + tests/e2e/bedrock/bedrock_test.go | 29 +++- tests/integration/bedrock_test.go | 25 +++- 15 files changed, 448 insertions(+), 58 deletions(-) create mode 100644 spinifex/gateway/bedrock/weights.go create mode 100644 spinifex/gateway/bedrock/weights_test.go diff --git a/spinifex/gateway/bedrock/catalog.go b/spinifex/gateway/bedrock/catalog.go index f416c7b5d..02cd40dc0 100644 --- a/spinifex/gateway/bedrock/catalog.go +++ b/spinifex/gateway/bedrock/catalog.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "log/slog" "strings" "github.com/aws/aws-sdk-go/aws" @@ -23,6 +24,10 @@ const ( // catalogEntry is one static catalog record. Model IDs mirror AWS exactly so // existing SDK code and configs are drop-in. +// +// MinVRAMMiB, InstanceType and VLLMArgs are the in-tree serving spec for +// self-host entries; the actual staged weights artifact is deployment-local +// state resolved from KV instead (see WeightsResolver). type catalogEntry struct { ModelID string ModelName string @@ -33,20 +38,29 @@ type catalogEntry struct { ResponseStreamingSupported bool InferenceTypesSupported []string CustomizationsSupported []string + MinVRAMMiB int // self-host only; 0 for provider entries + InstanceType string // self-host only; system-instance type a launcher boots + VLLMArgs []string // self-host only; extra vLLM server args this model needs } // catalog is the static Phase-1 model set: one self-hosted open model and one // Anthropic-direct model. Later phases extend this list. var catalog = []catalogEntry{ { - ModelID: "meta.llama3-70b-instruct-v1:0", - ModelName: "Llama 3 70B Instruct", + ModelID: "meta.llama3-2-1b-instruct-v1:0", + ModelName: "Llama 3.2 1B Instruct", ProviderName: "Meta", Provider: tierSelfHost, InputModalities: []string{"TEXT"}, OutputModalities: []string{"TEXT"}, ResponseStreamingSupported: false, InferenceTypesSupported: []string{"ON_DEMAND"}, + // MinVRAMMiB is the admission-gate floor; --gpu-memory-utilization + // caps vLLM's own pool to roughly the same figure (8188 MiB * 0.6 ≈ + // 4913 MiB) so the two stay consistent rather than drifting apart. + MinVRAMMiB: 5120, + InstanceType: "g5.xlarge", + VLLMArgs: []string{"--dtype=bfloat16", "--max-model-len=8192", "--gpu-memory-utilization=0.6"}, }, { ModelID: "anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -67,13 +81,23 @@ type CredentialResolver interface { Resolve(ctx context.Context, accountID, vendor string) (key string, ok bool, err error) } -// tieredCatalog returns the catalog entries advertised to accountID: self-host -// entries always, provider entries only when resolver finds a usable credential. -func tieredCatalog(ctx context.Context, accountID string, resolver CredentialResolver) []catalogEntry { +// tieredCatalog returns the catalog entries advertised to accountID: +// self-host entries only where a weights snapshot resolves, provider entries +// only when resolver finds a usable credential. A resolve error (as opposed +// to a clean not-found) is an internal fault, not a servability verdict, and +// aborts the whole list rather than silently thinning it out. +func tieredCatalog(ctx context.Context, accountID string, resolver CredentialResolver) ([]catalogEntry, error) { var out []catalogEntry for _, entry := range catalog { if entry.Provider == tierSelfHost { - out = append(out, entry) + _, resolvable, err := currentWeightsResolver().Resolve(ctx, entry.ModelID) + if err != nil { + slog.Error("bedrock: weights resolve failed", "model", entry.ModelID, "err", err) + return nil, fmt.Errorf("resolve weights for %s: %w", entry.ModelID, err) + } + if resolvable { + out = append(out, entry) + } continue } vendor, ok := strings.CutPrefix(entry.Provider, providerPrefix) @@ -86,7 +110,7 @@ func tieredCatalog(ctx context.Context, accountID string, resolver CredentialRes } out = append(out, entry) } - return out + return out, nil } func (e catalogEntry) toSummary() *bedrock.FoundationModelSummary { @@ -133,9 +157,13 @@ func lookupCatalogEntry(modelID string) (catalogEntry, bool) { } // ListFoundationModels returns the deployment-tiered catalog: self-host -// entries always, provider entries only where a credential resolves. +// entries where a weights snapshot resolves, provider entries where a +// credential resolves. func ListFoundationModels(ctx context.Context, accountID string, resolver CredentialResolver, _ *bedrock.ListFoundationModelsInput) (*bedrock.ListFoundationModelsOutput, error) { - entries := tieredCatalog(ctx, accountID, resolver) + entries, err := tieredCatalog(ctx, accountID, resolver) + if err != nil { + return nil, err + } summaries := make([]*bedrock.FoundationModelSummary, 0, len(entries)) for _, entry := range entries { summaries = append(summaries, entry.toSummary()) @@ -143,12 +171,24 @@ func ListFoundationModels(ctx context.Context, accountID string, resolver Creden return &bedrock.ListFoundationModelsOutput{ModelSummaries: summaries}, nil } -// GetFoundationModel looks up a single model by exact modelId, independent of -// tiering. Unknown models return ResourceNotFoundException. -func GetFoundationModel(_ context.Context, _ string, modelID string) (*bedrock.GetFoundationModelOutput, error) { +// GetFoundationModel looks up a single model by exact modelId. Unknown +// models, and self-host models with no resolvable weights snapshot, return +// ResourceNotFoundException. A weights resolve error is an internal fault, +// not a not-found verdict, and is surfaced (and logged) as such. +func GetFoundationModel(ctx context.Context, _ string, modelID string) (*bedrock.GetFoundationModelOutput, error) { entry, ok := lookupCatalogEntry(modelID) if !ok { return nil, errors.New(awserrors.ErrorResourceNotFoundException) } + if entry.Provider == tierSelfHost { + _, resolvable, err := currentWeightsResolver().Resolve(ctx, entry.ModelID) + if err != nil { + slog.Error("bedrock: weights resolve failed", "model", entry.ModelID, "err", err) + return nil, fmt.Errorf("resolve weights for %s: %w", entry.ModelID, err) + } + if !resolvable { + return nil, errors.New(awserrors.ErrorResourceNotFoundException) + } + } return &bedrock.GetFoundationModelOutput{ModelDetails: entry.toDetails()}, nil } diff --git a/spinifex/gateway/bedrock/catalog_test.go b/spinifex/gateway/bedrock/catalog_test.go index f3a9d11df..c5174624b 100644 --- a/spinifex/gateway/bedrock/catalog_test.go +++ b/spinifex/gateway/bedrock/catalog_test.go @@ -2,6 +2,7 @@ package gateway_bedrock import ( "context" + "errors" "testing" "github.com/aws/aws-sdk-go/service/bedrock" @@ -22,6 +23,37 @@ func (s stubResolver) Resolve(_ context.Context, _, vendor string) (string, bool return "stub-key", true, nil } +// stubWeightsResolver reports a resolvable weights snapshot only for model +// IDs in ok. +type stubWeightsResolver struct { + ok map[string]bool +} + +func (s stubWeightsResolver) Resolve(_ context.Context, modelID string) (string, bool, error) { + if !s.ok[modelID] { + return "", false, nil + } + return "snap-stub", true, nil +} + +// errWeightsResolver simulates a broken resolver (e.g. a JetStream outage): +// every Resolve call fails rather than cleanly reporting not-found. +type errWeightsResolver struct{} + +func (errWeightsResolver) Resolve(_ context.Context, _ string) (string, bool, error) { + return "", false, errors.New("kv unavailable") +} + +// withWeightsResolver installs r as the package-level weights resolver for +// the duration of the test, restoring the no-op default on cleanup — +// ListFoundationModels and GetFoundationModel read it via +// currentWeightsResolver rather than a parameter. +func withWeightsResolver(t *testing.T, r WeightsResolver) { + t.Helper() + SetWeightsResolver(r) + t.Cleanup(func() { SetWeightsResolver(nil) }) +} + func modelIDs(out *bedrock.ListFoundationModelsOutput) []string { ids := make([]string, 0, len(out.ModelSummaries)) for _, m := range out.ModelSummaries { @@ -30,29 +62,47 @@ func modelIDs(out *bedrock.ListFoundationModelsOutput) []string { return ids } -func TestListFoundationModels_SelfHostAlwaysIncluded(t *testing.T) { +func TestListFoundationModels_SelfHostIncludedWhenWeightsResolve(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{"meta.llama3-2-1b-instruct-v1:0": true}}) out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) require.NoError(t, err) - assert.Contains(t, modelIDs(out), "meta.llama3-70b-instruct-v1:0") + assert.Contains(t, modelIDs(out), "meta.llama3-2-1b-instruct-v1:0") +} + +func TestListFoundationModels_SelfHostExcludedWhenWeightsUnresolvable(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{}}) + out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) + require.NoError(t, err) + assert.NotContains(t, modelIDs(out), "meta.llama3-2-1b-instruct-v1:0") } func TestListFoundationModels_ProviderIncludedWhenResolvable(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{}}) out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{"anthropic": true}}, &bedrock.ListFoundationModelsInput{}) require.NoError(t, err) assert.Contains(t, modelIDs(out), "anthropic.claude-3-5-sonnet-20240620-v1:0") } func TestListFoundationModels_ProviderExcludedWhenUnresolvable(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{}}) out, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) require.NoError(t, err) assert.NotContains(t, modelIDs(out), "anthropic.claude-3-5-sonnet-20240620-v1:0") } -func TestGetFoundationModel_KnownModel(t *testing.T) { - out, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0") +func TestGetFoundationModel_KnownModelWithResolvableWeights(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{"meta.llama3-2-1b-instruct-v1:0": true}}) + out, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0") require.NoError(t, err) require.NotNil(t, out.ModelDetails) - assert.Equal(t, "meta.llama3-70b-instruct-v1:0", *out.ModelDetails.ModelId) + assert.Equal(t, "meta.llama3-2-1b-instruct-v1:0", *out.ModelDetails.ModelId) +} + +func TestGetFoundationModel_SelfHostWithUnresolvableWeightsReturnsNotFound(t *testing.T) { + withWeightsResolver(t, stubWeightsResolver{ok: map[string]bool{}}) + _, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0") + require.Error(t, err) + assert.Equal(t, "ResourceNotFoundException", err.Error()) } func TestGetFoundationModel_UnknownModelReturnsNotFound(t *testing.T) { @@ -60,3 +110,25 @@ func TestGetFoundationModel_UnknownModelReturnsNotFound(t *testing.T) { require.Error(t, err) assert.Equal(t, "ResourceNotFoundException", err.Error()) } + +func TestGetFoundationModel_WeightsResolveErrorIsNotResourceNotFound(t *testing.T) { + withWeightsResolver(t, errWeightsResolver{}) + _, err := GetFoundationModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0") + require.Error(t, err) + assert.NotEqual(t, "ResourceNotFoundException", err.Error()) +} + +func TestListFoundationModels_WeightsResolveErrorPropagates(t *testing.T) { + withWeightsResolver(t, errWeightsResolver{}) + _, err := ListFoundationModels(context.Background(), "000000000001", stubResolver{ok: map[string]bool{}}, &bedrock.ListFoundationModelsInput{}) + require.Error(t, err) + assert.NotEqual(t, "ResourceNotFoundException", err.Error()) +} + +func TestCatalog_SelfHostEntryCarriesServingSpec(t *testing.T) { + entry, ok := lookupCatalogEntry("meta.llama3-2-1b-instruct-v1:0") + require.True(t, ok) + assert.Positive(t, entry.MinVRAMMiB) + assert.NotEmpty(t, entry.InstanceType) + assert.NotEmpty(t, entry.VLLMArgs) +} diff --git a/spinifex/gateway/bedrock/contract_test.go b/spinifex/gateway/bedrock/contract_test.go index d3abd903e..dc079de61 100644 --- a/spinifex/gateway/bedrock/contract_test.go +++ b/spinifex/gateway/bedrock/contract_test.go @@ -44,7 +44,7 @@ func TestContract_ConverseStream_SelfHost_RealSDKConsumerDecodesFrames(t *testin })) defer vllmStub.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" endpoints := NewStaticEndpointResolver(map[string]string{modelID: vllmStub.URL}) svc := contractSession(t, func(w http.ResponseWriter, r *http.Request) { @@ -188,7 +188,7 @@ func TestContract_InvokeModelWithResponseStream_SelfHost_RealSDKConsumerDecodesF })) defer llamaStub.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" endpoints := NewStaticEndpointResolver(map[string]string{modelID: llamaStub.URL}) svc := contractSession(t, func(w http.ResponseWriter, r *http.Request) { diff --git a/spinifex/gateway/bedrock/converse_stream_test.go b/spinifex/gateway/bedrock/converse_stream_test.go index 9468555fa..0afb8d927 100644 --- a/spinifex/gateway/bedrock/converse_stream_test.go +++ b/spinifex/gateway/bedrock/converse_stream_test.go @@ -141,7 +141,7 @@ func TestConverseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing. func TestConverseStream_MalformedBodyReturnsValidationException(t *testing.T) { rec := httptest.NewRecorder() - err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-70b-instruct-v1:0", []byte("{not-json"), nil, nil) + err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte("{not-json"), nil, nil) require.Error(t, err) assert.Equal(t, awserrors.ErrorValidationException, err.Error()) } @@ -154,7 +154,7 @@ func TestConverseStream_SelfHostHappyPath_WritesFramedTaxonomy(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" rec := httptest.NewRecorder() body, err := json.Marshal(&bedrockruntime.ConverseStreamInput{ Messages: []*bedrockruntime.Message{ @@ -188,7 +188,7 @@ func TestConverseStream_NonFlusherWriter_ReturnsPreHeaderError(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" w := &nonFlusherWriter{} body, err := json.Marshal(&bedrockruntime.ConverseStreamInput{ Messages: []*bedrockruntime.Message{ diff --git a/spinifex/gateway/bedrock/invoke_llama_test.go b/spinifex/gateway/bedrock/invoke_llama_test.go index c8a6e5394..9d5b13c08 100644 --- a/spinifex/gateway/bedrock/invoke_llama_test.go +++ b/spinifex/gateway/bedrock/invoke_llama_test.go @@ -29,7 +29,7 @@ func TestLlamaInvokeAdapter_InvokeModel_MapsRequestAndResponse(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) a.httpClient = ts.Client() @@ -66,7 +66,7 @@ func TestLlamaInvokeAdapter_InvokeModel_MapsLengthFinishReason(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) a.httpClient = ts.Client() @@ -86,7 +86,7 @@ func TestLlamaInvokeAdapter_InvokeModel_EmptyChoicesReturnsModelError(t *testing })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) a.httpClient = ts.Client() @@ -98,7 +98,7 @@ func TestLlamaInvokeAdapter_InvokeModel_EmptyChoicesReturnsModelError(t *testing func TestLlamaInvokeAdapter_InvokeModel_UnresolvedEndpointReturnsModelNotReady(t *testing.T) { a := newLlamaInvokeAdapter(NewStaticEndpointResolver(nil)) - _, _, err := a.InvokeModel(context.Background(), "meta.llama3-70b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) + _, _, err := a.InvokeModel(context.Background(), "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } @@ -147,7 +147,7 @@ func TestLlamaInvokeAdapter_InvokeModelWithResponseStream_TranslatesChunks(t *te })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) a.httpClient = ts.Client() @@ -178,7 +178,7 @@ func TestLlamaInvokeAdapter_InvokeModelWithResponseStream_TranslatesChunks(t *te func TestLlamaInvokeAdapter_InvokeModelWithResponseStream_UnresolvedEndpointReturnsModelNotReady(t *testing.T) { a := newLlamaInvokeAdapter(NewStaticEndpointResolver(nil)) - _, err := a.InvokeModelWithResponseStream(context.Background(), "meta.llama3-70b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) + _, err := a.InvokeModelWithResponseStream(context.Background(), "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } @@ -191,7 +191,7 @@ func TestLlamaInvokeStreamSource_MidStreamDecodeErrorSurfacesAsStreamFault(t *te })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) a.httpClient = ts.Client() diff --git a/spinifex/gateway/bedrock/invoke_stream_test.go b/spinifex/gateway/bedrock/invoke_stream_test.go index 182e85cbb..1e23094f8 100644 --- a/spinifex/gateway/bedrock/invoke_stream_test.go +++ b/spinifex/gateway/bedrock/invoke_stream_test.go @@ -84,7 +84,7 @@ func TestInvokeModelWithResponseStream_SelfHostHappyPath_WritesChunkFrames(t *te })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" rec := httptest.NewRecorder() body := []byte(`{"prompt":"hello","max_gen_len":128}`) @@ -109,7 +109,7 @@ func TestInvokeModelWithResponseStream_NonFlusherWriter_ReturnsPreHeaderError(t })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" w := &nonFlusherWriter{} body := []byte(`{"prompt":"hello"}`) diff --git a/spinifex/gateway/bedrock/invoke_test.go b/spinifex/gateway/bedrock/invoke_test.go index b279fff97..756c3e8d4 100644 --- a/spinifex/gateway/bedrock/invoke_test.go +++ b/spinifex/gateway/bedrock/invoke_test.go @@ -23,7 +23,7 @@ func TestInvokeRouter_SelfHostSuccess(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) respBody, contentType, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`)) @@ -51,7 +51,7 @@ func TestInvokeRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) { func TestInvokeRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) { rt := NewInvokeRouter(nil, nil) - _, _, err := rt.InvokeModel(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) + _, _, err := rt.InvokeModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } @@ -67,7 +67,7 @@ func TestInvokeModel_PackageEntryPoint_SelfHostSuccess(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" respBody, contentType, err := InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) require.NoError(t, err) assert.Equal(t, "application/json", contentType) @@ -91,7 +91,7 @@ func TestInvokeStreamRouter_SelfHostSuccess(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) src, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`)) @@ -118,7 +118,7 @@ func TestInvokeStreamRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing. func TestInvokeStreamRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) { rt := NewInvokeStreamRouter(nil, nil) - _, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) + _, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`)) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } diff --git a/spinifex/gateway/bedrock/router_test.go b/spinifex/gateway/bedrock/router_test.go index 6cee3c9e2..9d6343033 100644 --- a/spinifex/gateway/bedrock/router_test.go +++ b/spinifex/gateway/bedrock/router_test.go @@ -45,7 +45,7 @@ func TestRouter_Converse_SelfHostSuccess(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) out, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput()) @@ -79,7 +79,7 @@ func TestConverse_PackageEntryPoint_SelfHostSuccess(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) require.NoError(t, err) require.NotNil(t, out.Output.Message) @@ -108,7 +108,7 @@ func TestRouter_ConverseStream_SelfHostSuccess(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) src, err := rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput()) @@ -135,7 +135,7 @@ func TestRouter_ConverseStream_AnthropicNoCredentialReturnsAccessDenied(t *testi func TestRouter_ConverseStream_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) { rt := NewRouter(nil, nil) - _, err := rt.ConverseStream(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", converseStreamInput()) + _, err := rt.ConverseStream(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", converseStreamInput()) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } @@ -146,7 +146,7 @@ func TestNewRouter_NilArgumentsFallBackToNoops(t *testing.T) { // Self-host model with no endpoint resolver configured resolves nothing, // so it must report ModelNotReady rather than panic on a nil resolver. - _, err := rt.Converse(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", converseInput()) + _, err := rt.Converse(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", converseInput()) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } diff --git a/spinifex/gateway/bedrock/vllm_test.go b/spinifex/gateway/bedrock/vllm_test.go index 7d3ca90a9..22bba34ed 100644 --- a/spinifex/gateway/bedrock/vllm_test.go +++ b/spinifex/gateway/bedrock/vllm_test.go @@ -31,7 +31,7 @@ func TestVLLMProvider_Converse_MapsResponse(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) p.httpClient = ts.Client() @@ -70,7 +70,7 @@ func TestVLLMProvider_Converse_UnresolvedEndpointReturnsModelNotReady(t *testing }, } - _, err := p.Converse(context.Background(), "meta.llama3-70b-instruct-v1:0", input) + _, err := p.Converse(context.Background(), "meta.llama3-2-1b-instruct-v1:0", input) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } @@ -123,7 +123,7 @@ func TestVLLMProvider_ConverseStream_MapsSSEToTaxonomy(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) p.httpClient = ts.Client() @@ -175,7 +175,7 @@ func TestVLLMProvider_ConverseStream_UnresolvedEndpointReturnsModelNotReady(t *t }, } - _, err := p.ConverseStream(context.Background(), "meta.llama3-70b-instruct-v1:0", input) + _, err := p.ConverseStream(context.Background(), "meta.llama3-2-1b-instruct-v1:0", input) require.Error(t, err) assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error()) } @@ -187,7 +187,7 @@ func TestVLLMProvider_ConverseStream_UpstreamNon2xxMapsStatus(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) p.httpClient = ts.Client() @@ -215,7 +215,7 @@ func TestVLLMProvider_ConverseStream_UsageWithoutFinishReasonKeepsOrder(t *testi })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) p.httpClient = ts.Client() @@ -247,7 +247,7 @@ func TestVLLMProvider_ConverseStream_EmptyStreamIsWellFormed(t *testing.T) { })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) p.httpClient = ts.Client() @@ -279,7 +279,7 @@ func TestVLLMConverseStreamSource_MidStreamDecodeErrorSurfacesAsStreamFault(t *t })) defer ts.Close() - modelID := "meta.llama3-70b-instruct-v1:0" + modelID := "meta.llama3-2-1b-instruct-v1:0" p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL})) p.httpClient = ts.Client() diff --git a/spinifex/gateway/bedrock/weights.go b/spinifex/gateway/bedrock/weights.go new file mode 100644 index 000000000..d508eba23 --- /dev/null +++ b/spinifex/gateway/bedrock/weights.go @@ -0,0 +1,138 @@ +package gateway_bedrock + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "sync" + + "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/nats-io/nats.go/jetstream" +) + +// bedrockWeightsBucket is the cluster-replicated KV bucket holding, per +// model, the snapshot ID of this deployment's staged weights volume. Two +// deployments of the same catalog can stage different snapshots, or none. +const bedrockWeightsBucket = "bedrock-weights" + +// bedrockWeightsHistory keeps one revision; a re-stage overwrites in place. +const bedrockWeightsHistory = 1 + +// weightsKey returns the KV key for modelID's staged-weights snapshot. +// Model IDs contain ':' (e.g. "meta.llama3-2-1b-instruct-v1:0"), which NATS +// rejects in a KV key, so the segment is base64url-encoded. +func weightsKey(modelID string) string { + return base64.RawURLEncoding.EncodeToString([]byte(modelID)) +} + +// WeightsResolver resolves a self-hosted model's deployment-specific weights +// snapshot ID. A model with no resolvable snapshot has nothing to serve it +// with and must not be advertised — see tieredCatalog and GetFoundationModel. +type WeightsResolver interface { + Resolve(ctx context.Context, modelID string) (snapshotID string, ok bool, err error) +} + +// WeightsStore resolves per-model weights snapshot IDs from the +// bedrock-weights JetStream KV bucket. +type WeightsStore struct { + js jetstream.JetStream + replicas int + + mu sync.Mutex + kv jetstream.KeyValue +} + +var _ WeightsResolver = (*WeightsStore)(nil) + +// NewWeightsStore constructs a WeightsStore over the cluster's JetStream +// client, replicated across replicas nodes. +func NewWeightsStore(js jetstream.JetStream, replicas int) *WeightsStore { + return &WeightsStore{js: js, replicas: replicas} +} + +// bucket lazily opens (or creates) the cluster-replicated bedrock-weights KV +// bucket, caching the handle for subsequent calls, mirroring +// CredentialStore.bucket. +func (s *WeightsStore) bucket(ctx context.Context) (jetstream.KeyValue, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.kv != nil { + return s.kv, nil + } + kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockWeightsBucket, bedrockWeightsHistory, s.replicas) + if err != nil { + return nil, err + } + s.kv = kv + return kv, nil +} + +// Resolve returns modelID's staged weights snapshot ID, if one has been set. +func (s *WeightsStore) Resolve(ctx context.Context, modelID string) (string, bool, error) { + kv, err := s.bucket(ctx) + if err != nil { + return "", false, err + } + entry, err := kv.Get(ctx, weightsKey(modelID)) + switch { + case err == nil: + return string(entry.Value()), true, nil + case errors.Is(err, jetstream.ErrKeyNotFound): + return "", false, nil + default: + return "", false, fmt.Errorf("kv get weights snapshot for %s: %w", modelID, err) + } +} + +// PutWeights records snapshotID as modelID's staged weights artifact. +func (s *WeightsStore) PutWeights(ctx context.Context, modelID, snapshotID string) error { + kv, err := s.bucket(ctx) + if err != nil { + return err + } + if _, err := kv.Put(ctx, weightsKey(modelID), []byte(snapshotID)); err != nil { + return fmt.Errorf("kv put weights snapshot for %s: %w", modelID, err) + } + return nil +} + +// noopWeightsResolver resolves no snapshot for any model. +type noopWeightsResolver struct{} + +var _ WeightsResolver = (*noopWeightsResolver)(nil) + +func (noopWeightsResolver) Resolve(_ context.Context, _ string) (string, bool, error) { + return "", false, nil +} + +// NoopWeightsResolver resolves no weights for any model: the unconfigured +// direction is "no self-host model is servable", matching how +// NoopCredentialResolver resolves no provider credentials. +var NoopWeightsResolver WeightsResolver = noopWeightsResolver{} + +// weightsResolverMu guards weightsResolver, process-wide state set once at +// service start via SetWeightsResolver. +var ( + weightsResolverMu sync.RWMutex + weightsResolver WeightsResolver = NoopWeightsResolver +) + +// SetWeightsResolver installs the resolver tieredCatalog and +// GetFoundationModel gate self-host entries on. A nil resolver restores the +// no-op default. +func SetWeightsResolver(r WeightsResolver) { + weightsResolverMu.Lock() + defer weightsResolverMu.Unlock() + if r == nil { + r = NoopWeightsResolver + } + weightsResolver = r +} + +// currentWeightsResolver returns the resolver installed by SetWeightsResolver. +func currentWeightsResolver() WeightsResolver { + weightsResolverMu.RLock() + defer weightsResolverMu.RUnlock() + return weightsResolver +} diff --git a/spinifex/gateway/bedrock/weights_test.go b/spinifex/gateway/bedrock/weights_test.go new file mode 100644 index 000000000..ebb391390 --- /dev/null +++ b/spinifex/gateway/bedrock/weights_test.go @@ -0,0 +1,77 @@ +package gateway_bedrock + +import ( + "context" + "testing" + + "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNoopWeightsResolver_ResolvesNothing(t *testing.T) { + snapshotID, ok, err := NoopWeightsResolver.Resolve(context.Background(), "meta.llama3-2-1b-instruct-v1:0") + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, snapshotID) +} + +// TestWeightsStore_ResolveMiss covers a KV miss on a model with no staged +// snapshot: ("", false, nil), not an error. +func TestWeightsStore_ResolveMiss(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + store := NewWeightsStore(testutil.NewJetStream(t, nc), 1) + + snapshotID, ok, err := store.Resolve(context.Background(), "meta.llama3-2-1b-instruct-v1:0") + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, snapshotID) +} + +// TestWeightsStore_PutAndResolve_KV exercises the real JetStream KV path: +// bucket (lazy create), weightsKey, PutWeights, and the KV-hit branch of +// Resolve. +func TestWeightsStore_PutAndResolve_KV(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + store := NewWeightsStore(testutil.NewJetStream(t, nc), 1) + + ctx := context.Background() + require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "snap-0001")) + + snapshotID, ok, err := store.Resolve(ctx, "meta.llama3-2-1b-instruct-v1:0") + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, "snap-0001", snapshotID) +} + +// TestWeightsStore_PutOverwritesPrevious covers a re-stage: PutWeights +// overwrites the prior snapshot ID in place rather than erroring or +// accumulating history. +func TestWeightsStore_PutOverwritesPrevious(t *testing.T) { + _, nc, _ := testutil.StartTestJetStream(t) + store := NewWeightsStore(testutil.NewJetStream(t, nc), 1) + + ctx := context.Background() + require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "snap-0001")) + require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "snap-0002")) + + snapshotID, ok, err := store.Resolve(ctx, "meta.llama3-2-1b-instruct-v1:0") + require.NoError(t, err) + assert.True(t, ok) + assert.Equal(t, "snap-0002", snapshotID) +} + +func TestWeightsKey(t *testing.T) { + // base64url(RawURLEncoding) of "meta.llama3-2-1b-instruct-v1:0". + assert.Equal(t, "bWV0YS5sbGFtYTMtMi0xYi1pbnN0cnVjdC12MTow", weightsKey("meta.llama3-2-1b-instruct-v1:0")) +} + +func TestSetWeightsResolver_NilRestoresNoop(t *testing.T) { + SetWeightsResolver(stubWeightsResolver{ok: map[string]bool{"x": true}}) + SetWeightsResolver(nil) + t.Cleanup(func() { SetWeightsResolver(nil) }) + + _, ok, err := currentWeightsResolver().Resolve(context.Background(), "x") + require.NoError(t, err) + assert.False(t, ok) +} diff --git a/spinifex/gateway/bedrock_request_test.go b/spinifex/gateway/bedrock_request_test.go index 581fc4380..7663b11e4 100644 --- a/spinifex/gateway/bedrock_request_test.go +++ b/spinifex/gateway/bedrock_request_test.go @@ -11,13 +11,14 @@ import ( "github.com/aws/aws-sdk-go/service/bedrock" "github.com/aws/aws-sdk-go/service/bedrockruntime" "github.com/mulgadc/spinifex/spinifex/awserrors" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" "github.com/mulgadc/spinifex/spinifex/testutil" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) const bedrockTestAccount = "123456789012" -const bedrockTestLlamaModelID = "meta.llama3-70b-instruct-v1:0" +const bedrockTestLlamaModelID = "meta.llama3-2-1b-instruct-v1:0" // newVLLMStub stands up an httptest server that answers the OpenAI // chat-completions wire, mirroring gateway/bedrock/vllm_test.go's stub. @@ -52,13 +53,30 @@ func newLlamaCompletionsStub(t *testing.T) *httptest.Server { return ts } +// bedrockRequestWeightsStub resolves a weights snapshot only for +// bedrockTestLlamaModelID, standing in for the KV override +// ListFoundationModels/GetFoundationModel gate self-host entries on. +type bedrockRequestWeightsStub struct{} + +func (bedrockRequestWeightsStub) Resolve(_ context.Context, modelID string) (string, bool, error) { + if modelID != bedrockTestLlamaModelID { + return "", false, nil + } + return "snap-stub", true, nil +} + // newBedrockRequestGateway builds a GatewayConfig with a real NATS connection // (satisfying Bedrock_Request/BedrockRuntime_Request's nil-check only — no // NATS subject handling is required for these routes), no IAMService (so // checkPolicy is a no-op), and BedrockEndpoints pinned at the given vLLM stub. +// ListFoundationModels/GetFoundationModel read the weights resolver through a +// package-level setter rather than a GatewayConfig field, so it is installed +// here and reset on cleanup. func newBedrockRequestGateway(t *testing.T, vllmURL string) *GatewayConfig { t.Helper() _, nc, _ := testutil.StartTestJetStream(t) + gateway_bedrock.SetWeightsResolver(bedrockRequestWeightsStub{}) + t.Cleanup(func() { gateway_bedrock.SetWeightsResolver(nil) }) return &GatewayConfig{ NATSConn: nc, DisableLogging: true, diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index e765e933f..67915091f 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -301,6 +301,13 @@ func launchService(config *config.ClusterConfig) error { } bedrockCredentials := gateway_bedrock.NewCredentialStore(js, masterKey, len(config.Nodes), bedrockPlatformDefaults) + // Bedrock self-host weights: a model's serving spec (VRAM, instance type, + // vLLM args) ships in-tree, but which staged snapshot serves it is + // deployment-local state. tieredCatalog/GetFoundationModel read this + // through the package-level resolver rather than a parameter, since they + // are called from gateway/bedrock.go's fixed-arity route table. + gateway_bedrock.SetWeightsResolver(gateway_bedrock.NewWeightsStore(js, len(config.Nodes))) + // Bedrock self-host endpoints: Phase 1 models are pinned, so their // OpenAI-compatible base URLs come from static config. OCHRE_VLLM_ENDPOINTS // is a comma-separated list of modelId=baseURL pairs. diff --git a/tests/e2e/bedrock/bedrock_test.go b/tests/e2e/bedrock/bedrock_test.go index 1be400320..41076f305 100644 --- a/tests/e2e/bedrock/bedrock_test.go +++ b/tests/e2e/bedrock/bedrock_test.go @@ -17,23 +17,25 @@ import ( // Catalog model IDs mirror the static Phase-1 catalog in gateway_bedrock. const ( - modelSelfHostLlama = "meta.llama3-70b-instruct-v1:0" + modelSelfHostLlama = "meta.llama3-2-1b-instruct-v1:0" modelAnthropic = "anthropic.claude-3-5-sonnet-20240620-v1:0" modelBogus = "does.not-exist-v1:0" ) // TestGetFoundationModel round-trips a known model and asserts -// ResourceNotFoundException for an unknown one. +// ResourceNotFoundException for an unknown one. Uses the Anthropic entry: it +// resolves unconditionally, unlike the self-host entry, which is gated on a +// weights snapshot being staged — see TestGetFoundationModelSelfHost. func TestGetFoundationModel(t *testing.T) { f := requireBedrockFixture(t) t.Run("known model", func(t *testing.T) { out, err := f.AWS.Bedrock.GetFoundationModel(&bedrock.GetFoundationModelInput{ - ModelIdentifier: aws.String(modelSelfHostLlama), + ModelIdentifier: aws.String(modelAnthropic), }) - require.NoError(t, err, "get-foundation-model %s", modelSelfHostLlama) + require.NoError(t, err, "get-foundation-model %s", modelAnthropic) require.NotNil(t, out.ModelDetails, "empty ModelDetails") - assert.Equal(t, modelSelfHostLlama, aws.StringValue(out.ModelDetails.ModelId)) + assert.Equal(t, modelAnthropic, aws.StringValue(out.ModelDetails.ModelId)) }) t.Run("unknown model", func(t *testing.T) { @@ -46,6 +48,23 @@ func TestGetFoundationModel(t *testing.T) { }) } +// TestGetFoundationModelSelfHost describes the self-host model directly. +// Gated behind OCHRE_E2E_SELFHOST=1: besides a GPU-backed vLLM endpoint, it +// needs the model's weights snapshot staged in the bedrock-weights KV bucket. +func TestGetFoundationModelSelfHost(t *testing.T) { + if os.Getenv("OCHRE_E2E_SELFHOST") != "1" { + t.Skip("OCHRE_E2E_SELFHOST!=1; skipping self-host GetFoundationModel") + } + f := requireBedrockFixture(t) + + out, err := f.AWS.Bedrock.GetFoundationModel(&bedrock.GetFoundationModelInput{ + ModelIdentifier: aws.String(modelSelfHostLlama), + }) + require.NoError(t, err, "get-foundation-model %s", modelSelfHostLlama) + require.NotNil(t, out.ModelDetails, "empty ModelDetails") + assert.Equal(t, modelSelfHostLlama, aws.StringValue(out.ModelDetails.ModelId)) +} + // TestConverseSelfHost exercises the real vLLM Converse path. Gated behind // OCHRE_E2E_SELFHOST=1 — it needs a GPU-backed vLLM endpoint configured on the // gateway via OCHRE_VLLM_ENDPOINTS at boot. diff --git a/tests/integration/bedrock_test.go b/tests/integration/bedrock_test.go index 1acda1f10..09e483bd6 100644 --- a/tests/integration/bedrock_test.go +++ b/tests/integration/bedrock_test.go @@ -3,10 +3,12 @@ package integration import ( + "context" "testing" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/bedrock" + gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -15,17 +17,34 @@ import ( // gateway/bedrock/catalog.go — see tests/e2e/bedrock/bedrock_test.go, whose // TestListFoundationModels this ports. const ( - modelSelfHostLlama = "meta.llama3-70b-instruct-v1:0" + modelSelfHostLlama = "meta.llama3-2-1b-instruct-v1:0" modelAnthropic = "anthropic.claude-3-5-sonnet-20240620-v1:0" ) +// integrationWeightsStub resolves a weights snapshot only for +// modelSelfHostLlama, standing in for the KV override that gates a self-host +// catalog entry on having somewhere to serve it from. +type integrationWeightsStub struct{} + +func (integrationWeightsStub) Resolve(_ context.Context, modelID string) (string, bool, error) { + if modelID != modelSelfHostLlama { + return "", false, nil + } + return "snap-stub", true, nil +} + // TestListFoundationModels ports tests/e2e/bedrock/bedrock_test.go's // TestListFoundationModels. ListFoundationModels is pure gateway-side static // catalog logic (gateway/bedrock/catalog.go) with no NATS/daemon hop at all — // no guest, no daemon wiring needed — so it round-trips through the real // gateway HTTP/SigV4 path exactly like the live test, without a live cluster. func TestListFoundationModels(t *testing.T) { - t.Parallel() + // Not t.Parallel(): ListFoundationModels/GetFoundationModel read the + // weights resolver through a package-level setter (StartGateway builds a + // GatewayConfig directly, with no field for it), so this test owns that + // process-wide state for its duration. + gateway_bedrock.SetWeightsResolver(integrationWeightsStub{}) + t.Cleanup(func() { gateway_bedrock.SetWeightsResolver(nil) }) gw := StartGateway(t) bedrockCli := gw.BedrockClient(t) @@ -38,7 +57,7 @@ func TestListFoundationModels(t *testing.T) { ids[aws.StringValue(m.ModelId)] = m } - assert.Contains(t, ids, modelSelfHostLlama, "self-host model must always be advertised") + assert.Contains(t, ids, modelSelfHostLlama, "self-host model must be advertised when its weights resolve") // This tier configures no BedrockCredentials resolver, so the Anthropic // entry is never resolvable and must not be advertised — the mirror image