gsadmin/spinifex
1
0
Fork 0
mirror of https://github.com/mulgadc/spinifex.git synced 2026-10-21 19:56:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: convert the bedrock weights, prices and logging stores to kvstore (#895)

Browse Source
This commit is contained in:
Josh Siviour
2026-08-26 23:23:25 +10:00
committed by GitHub
parent e24fd42bdc
commit 70414154c5
6 changed files with 110 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spinifex/gateway/bedrock/logging.go
+18 -57
View File
@@ -9,7 +9,6 @@ import (
"log/slog"
"net/http"
"strings"
"sync"
"sync/atomic"
"time"
@@ -17,7 +16,7 @@ import (
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/kvutil"
"github.com/mulgadc/spinifex/spinifex/kvstore"
"github.com/mulgadc/spinifex/spinifex/objectstore"
"github.com/nats-io/nats.go/jetstream"
)
@@ -305,72 +304,38 @@ const bedrockLoggingConfigHistory = 1
// bedrock-logging-config JetStream KV bucket, and serves as the
// LoggingConfigReader StreamRecorder consults before including body text.
type LoggingConfigStore struct {
js jetstream.JetStream
replicas int
mu sync.Mutex
kv jetstream.KeyValue
store *kvstore.Store[LoggingConfig]
}
var _ LoggingConfigReader = (*LoggingConfigStore)(nil)
// NewLoggingConfigStore constructs a LoggingConfigStore.
func NewLoggingConfigStore(js jetstream.JetStream, replicas int) *LoggingConfigStore {
return &LoggingConfigStore{js: js, replicas: replicas}
}
// bucket lazily opens (or creates) the cluster-replicated
// bedrock-logging-config KV bucket, caching the handle for subsequent calls.
func (s *LoggingConfigStore) bucket(ctx context.Context) (jetstream.KeyValue, error) {
if s.js == nil {
return nil, errors.New("bedrock: logging config store has no JetStream client configured")
}
s.mu.Lock()
defer s.mu.Unlock()
if s.kv != nil {
return s.kv, nil
}
kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockLoggingConfigBucket, bedrockLoggingConfigHistory, s.replicas)
if err != nil {
return nil, err
}
s.kv = kv
return kv, nil
return &LoggingConfigStore{store: kvstore.New[LoggingConfig](js, kvstore.Config{
Name: bedrockLoggingConfigBucket,
History: bedrockLoggingConfigHistory,
Replicas: replicas,
Missing: "bedrock: logging config store has no JetStream client configured",
})}
}
// Get returns accountID's stored logging config, or (zero, false, nil) if
// the account has never configured one.
func (s *LoggingConfigStore) Get(ctx context.Context, accountID string) (LoggingConfig, bool, error) {
kv, err := s.bucket(ctx)
cfg, _, err := s.store.Get(ctx, accountID)
if errors.Is(err, kvstore.ErrNotFound) {
return LoggingConfig{}, false, nil
}
if err != nil {
return LoggingConfig{}, false, err
return LoggingConfig{}, false, fmt.Errorf("logging config for %s: %w", accountID, err)
}
entry, err := kv.Get(ctx, accountID)
if err != nil {
if errors.Is(err, jetstream.ErrKeyNotFound) {
return LoggingConfig{}, false, nil
}
return LoggingConfig{}, false, fmt.Errorf("kv get logging config for %s: %w", accountID, err)
}
var cfg LoggingConfig
if err := json.Unmarshal(entry.Value(), &cfg); err != nil {
return LoggingConfig{}, false, fmt.Errorf("decode logging config for %s: %w", accountID, err)
}
return cfg, true, nil
return *cfg, true, nil
}
// Put stores accountID's logging config, overwriting any existing one.
func (s *LoggingConfigStore) Put(ctx context.Context, accountID string, cfg LoggingConfig) error {
kv, err := s.bucket(ctx)
if err != nil {
return err
}
data, err := json.Marshal(cfg)
if err != nil {
return fmt.Errorf("encode logging config for %s: %w", accountID, err)
}
if _, err := kv.Put(ctx, accountID, data); err != nil {
return fmt.Errorf("kv put logging config for %s: %w", accountID, err)
if err := s.store.Set(ctx, accountID, &cfg); err != nil {
return fmt.Errorf("logging config for %s: %w", accountID, err)
}
return nil
}
@@ -378,12 +343,8 @@ func (s *LoggingConfigStore) Put(ctx context.Context, accountID string, cfg Logg
// Delete removes accountID's logging config, treating an already-absent
// config as success.
func (s *LoggingConfigStore) Delete(ctx context.Context, accountID string) error {
kv, err := s.bucket(ctx)
if err != nil {
return err
}
if err := kv.Delete(ctx, accountID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
return fmt.Errorf("kv delete logging config for %s: %w", accountID, err)
if err := s.store.Delete(ctx, accountID); err != nil {
return fmt.Errorf("logging config for %s: %w", accountID, err)
}
return nil
}
spinifex/gateway/bedrock/logging_test.go
+15
View File
@@ -289,3 +289,18 @@ func TestDeliveryConsumer_Run_DeliversMetadataOnlyRecordToS3(t *testing.T) {
assert.Empty(t, stored.InputText)
assert.Empty(t, stored.OutputText)
}
// TestLoggingConfigStore_RequiresJetStream covers a store constructed with no
// JetStream client: every accessor must report the misconfiguration rather
// than panic on a nil handle.
func TestLoggingConfigStore_RequiresJetStream(t *testing.T) {
store := NewLoggingConfigStore(nil, 1)
ctx := context.Background()
_, _, err := store.Get(ctx, "000000000001")
require.Error(t, err)
assert.Contains(t, err.Error(), "no JetStream client configured")
require.Error(t, store.Put(ctx, "000000000001", LoggingConfig{}))
require.Error(t, store.Delete(ctx, "000000000001"))
}
spinifex/gateway/bedrock/prices.go
+20 -58
View File
@@ -2,12 +2,10 @@ package gateway_bedrock
import (
"context"
"encoding/json"
"errors"
"fmt"
"sync"
"github.com/mulgadc/spinifex/spinifex/kvutil"
"github.com/mulgadc/spinifex/spinifex/kvstore"
"github.com/nats-io/nats.go/jetstream"
)
@@ -72,11 +70,7 @@ const bedrockPricesHistory = 1
// — both are per-model deployment data with an in-tree default plus a KV
// override (D4, D12).
type PriceStore struct {
js jetstream.JetStream
replicas int
mu sync.Mutex
kv jetstream.KeyValue
store *kvstore.Store[Price]
}
var _ PriceResolver = (*PriceStore)(nil)
@@ -84,62 +78,34 @@ var _ PriceResolver = (*PriceStore)(nil)
// NewPriceStore constructs a PriceStore over the cluster's JetStream client,
// replicated across replicas nodes.
func NewPriceStore(js jetstream.JetStream, replicas int) *PriceStore {
return &PriceStore{js: js, replicas: replicas}
return &PriceStore{store: kvstore.New[Price](js, kvstore.Config{
Name: bedrockPricesBucket,
History: bedrockPricesHistory,
Replicas: replicas,
Missing: "bedrock: price store has no JetStream client configured",
})}
}
// bucket lazily opens (or creates) the cluster-replicated bedrock-prices KV
// bucket, caching the handle for subsequent calls, mirroring
// WeightsStore.bucket.
func (s *PriceStore) 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, bedrockPricesBucket, bedrockPricesHistory, s.replicas)
if err != nil {
return nil, err
}
s.kv = kv
return kv, nil
}
// Resolve returns modelID's KV-overridden price, if one has been set.
// Resolve returns modelID's KV-overridden price, if one has been set. Errors
// name the model ID, since the key itself is base64url.
func (s *PriceStore) Resolve(ctx context.Context, modelID string) (Price, bool, error) {
kv, err := s.bucket(ctx)
if err != nil {
return Price{}, false, err
}
entry, err := kv.Get(ctx, weightsKey(modelID))
switch {
case err == nil:
var price Price
if err := json.Unmarshal(entry.Value(), &price); err != nil {
return Price{}, false, fmt.Errorf("decode price override for %s: %w", modelID, err)
}
return price, true, nil
case errors.Is(err, jetstream.ErrKeyNotFound):
price, _, err := s.store.Get(ctx, weightsKey(modelID))
if errors.Is(err, kvstore.ErrNotFound) {
return Price{}, false, nil
default:
return Price{}, false, fmt.Errorf("kv get price override for %s: %w", modelID, err)
}
if err != nil {
return Price{}, false, fmt.Errorf("price override for %s: %w", modelID, err)
}
return *price, true, nil
}
// PutPrice records price as modelID's overridden price, always Known
// (storing a KV override to mark a model unknown makes no sense — remove the
// override with DeletePrice instead to fall back to the in-tree default).
func (s *PriceStore) PutPrice(ctx context.Context, modelID string, price Price) error {
kv, err := s.bucket(ctx)
if err != nil {
return err
}
price.Known = true
data, err := json.Marshal(price)
if err != nil {
return fmt.Errorf("encode price override for %s: %w", modelID, err)
}
if _, err := kv.Put(ctx, weightsKey(modelID), data); err != nil {
return fmt.Errorf("kv put price override for %s: %w", modelID, err)
if err := s.store.Set(ctx, weightsKey(modelID), &price); err != nil {
return fmt.Errorf("price override for %s: %w", modelID, err)
}
return nil
}
@@ -148,12 +114,8 @@ func (s *PriceStore) PutPrice(ctx context.Context, modelID string, price Price)
// the catalog entry's in-tree default. Deleting an already-absent override is
// idempotent, not an error.
func (s *PriceStore) DeletePrice(ctx context.Context, modelID string) error {
kv, err := s.bucket(ctx)
if err != nil {
return err
}
if err := kv.Delete(ctx, weightsKey(modelID)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
return fmt.Errorf("kv delete price override for %s: %w", modelID, err)
if err := s.store.Delete(ctx, weightsKey(modelID)); err != nil {
return fmt.Errorf("price override for %s: %w", modelID, err)
}
return nil
}
spinifex/gateway/bedrock/prices_test.go
+15
View File
@@ -121,3 +121,18 @@ func TestPriceStore_PutGetDelete(t *testing.T) {
// Deleting an already-absent override is idempotent, not an error.
require.NoError(t, store.DeletePrice(ctx, priceTestProviderModelID))
}
// TestPriceStore_RequiresJetStream covers a store constructed with no
// JetStream client: every accessor must report the misconfiguration rather
// than panic on a nil handle.
func TestPriceStore_RequiresJetStream(t *testing.T) {
store := NewPriceStore(nil, 1)
ctx := context.Background()
_, _, err := store.Resolve(ctx, priceTestProviderModelID)
require.Error(t, err)
assert.Contains(t, err.Error(), "no JetStream client configured")
require.Error(t, store.PutPrice(ctx, priceTestProviderModelID, Price{}))
require.Error(t, store.DeletePrice(ctx, priceTestProviderModelID))
}
spinifex/gateway/bedrock/weights.go
+24 -56
View File
@@ -9,7 +9,7 @@ import (
"sort"
"sync"
"github.com/mulgadc/spinifex/spinifex/kvutil"
"github.com/mulgadc/spinifex/spinifex/kvstore"
"github.com/nats-io/nats.go/jetstream"
)
@@ -51,11 +51,7 @@ type WeightsResolver interface {
// WeightsStore resolves per-model weights records from the bedrock-weights
// JetStream KV bucket.
type WeightsStore struct {
js jetstream.JetStream
replicas int
mu sync.Mutex
kv jetstream.KeyValue
store *kvstore.Store[weightsRecord]
}
var _ WeightsResolver = (*WeightsStore)(nil)
@@ -63,47 +59,27 @@ 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
return &WeightsStore{store: kvstore.New[weightsRecord](js, kvstore.Config{
Name: bedrockWeightsBucket,
History: bedrockWeightsHistory,
Replicas: replicas,
Missing: "bedrock: weights store has no JetStream client configured",
})}
}
// getRecord reads and decodes modelID's weightsRecord. ok is false on a KV
// miss (not staged); a malformed record is reported as an error rather than
// silently treated as unstaged, since that would mask real corruption.
// Errors name the model ID, since the key itself is base64url.
func (s *WeightsStore) getRecord(ctx context.Context, modelID string) (weightsRecord, bool, error) {
kv, err := s.bucket(ctx)
if err != nil {
return weightsRecord{}, false, err
}
entry, err := kv.Get(ctx, weightsKey(modelID))
switch {
case err == nil:
var rec weightsRecord
if jsonErr := json.Unmarshal(entry.Value(), &rec); jsonErr != nil {
return weightsRecord{}, false, fmt.Errorf("decode weights record for %s: %w", modelID, jsonErr)
}
return rec, true, nil
case errors.Is(err, jetstream.ErrKeyNotFound):
rec, _, err := s.store.Get(ctx, weightsKey(modelID))
if errors.Is(err, kvstore.ErrNotFound) {
return weightsRecord{}, false, nil
default:
return weightsRecord{}, false, fmt.Errorf("kv get weights record for %s: %w", modelID, err)
}
if err != nil {
return weightsRecord{}, false, fmt.Errorf("weights record for %s: %w", modelID, err)
}
return *rec, true, nil
}
// Resolve returns modelID's staged weights snapshot ID, if one has been set.
@@ -138,16 +114,9 @@ func (s *WeightsStore) PutWeights(ctx context.Context, modelID, sourceURI, snaps
// the upstream hub commit SHA it was pulled at, if known (empty when 'stage'
// found no 'ochre-pull.json' manifest at the source prefix).
func (s *WeightsStore) PutWeightsWithRevision(ctx context.Context, modelID, sourceURI, snapshotID, sourceRevision string) error {
kv, err := s.bucket(ctx)
if err != nil {
return err
}
value, err := json.Marshal(weightsRecord{SnapshotID: snapshotID, SourceURI: sourceURI, SourceRevision: sourceRevision})
if err != nil {
return fmt.Errorf("encode weights record for %s: %w", modelID, err)
}
if _, err := kv.Put(ctx, weightsKey(modelID), value); err != nil {
return fmt.Errorf("kv put weights record for %s: %w", modelID, err)
rec := weightsRecord{SnapshotID: snapshotID, SourceURI: sourceURI, SourceRevision: sourceRevision}
if err := s.store.Set(ctx, weightsKey(modelID), &rec); err != nil {
return fmt.Errorf("weights record for %s: %w", modelID, err)
}
return nil
}
@@ -165,8 +134,11 @@ type WeightsEntry struct {
// ID, so 'ochre weights list' can show an operator what's staged, where it
// came from, and why a model is (or isn't) advertised via
// ListFoundationModels.
// It stays on the raw handle rather than Store.List for two reasons: the model
// ID is recoverable only from the key, which Store.List does not surface, and
// one unreadable record must not deny an operator the listing of every other.
func (s *WeightsStore) ListWeights(ctx context.Context) ([]WeightsEntry, error) {
kv, err := s.bucket(ctx)
kv, err := s.store.KV(ctx)
if err != nil {
return nil, err
}
@@ -203,12 +175,8 @@ func (s *WeightsStore) ListWeights(ctx context.Context) ([]WeightsEntry, error)
// volume, snapshot and source S3 objects are left intact — reclaiming that
// storage is a separate, explicit act.
func (s *WeightsStore) DeleteWeights(ctx context.Context, modelID string) error {
kv, err := s.bucket(ctx)
if err != nil {
return err
}
if err := kv.Delete(ctx, weightsKey(modelID)); err != nil {
return fmt.Errorf("kv delete weights record for %s: %w", modelID, err)
if err := s.store.Delete(ctx, weightsKey(modelID)); err != nil {
return fmt.Errorf("weights record for %s: %w", modelID, err)
}
return nil
}
spinifex/gateway/bedrock/weights_test.go
+18
View File
@@ -203,3 +203,21 @@ func TestSetWeightsResolver_NilRestoresNoop(t *testing.T) {
require.NoError(t, err)
assert.False(t, ok)
}
// TestWeightsStore_RequiresJetStream covers a store constructed with no
// JetStream client: every accessor must report the misconfiguration rather
// than panic on a nil handle.
func TestWeightsStore_RequiresJetStream(t *testing.T) {
store := NewWeightsStore(nil, 1)
ctx := context.Background()
_, _, err := store.Resolve(ctx, "meta.llama3-2-1b-instruct-v1:0")
require.Error(t, err)
assert.Contains(t, err.Error(), "no JetStream client configured")
require.Error(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/x/", "snap-0001"))
require.Error(t, store.DeleteWeights(ctx, "meta.llama3-2-1b-instruct-v1:0"))
_, err = store.ListWeights(ctx)
require.Error(t, err)
}