diff --git a/spinifex/gateway/bedrock/logging.go b/spinifex/gateway/bedrock/logging.go index a701d18be..fb3f9c6d5 100644 --- a/spinifex/gateway/bedrock/logging.go +++ b/spinifex/gateway/bedrock/logging.go @@ -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 } diff --git a/spinifex/gateway/bedrock/logging_test.go b/spinifex/gateway/bedrock/logging_test.go index 67c9cc7de..1d21b975f 100644 --- a/spinifex/gateway/bedrock/logging_test.go +++ b/spinifex/gateway/bedrock/logging_test.go @@ -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")) +} diff --git a/spinifex/gateway/bedrock/prices.go b/spinifex/gateway/bedrock/prices.go index d6094681f..7d74f0dd5 100644 --- a/spinifex/gateway/bedrock/prices.go +++ b/spinifex/gateway/bedrock/prices.go @@ -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 } diff --git a/spinifex/gateway/bedrock/prices_test.go b/spinifex/gateway/bedrock/prices_test.go index 7ac95a120..573a7f059 100644 --- a/spinifex/gateway/bedrock/prices_test.go +++ b/spinifex/gateway/bedrock/prices_test.go @@ -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)) +} diff --git a/spinifex/gateway/bedrock/weights.go b/spinifex/gateway/bedrock/weights.go index 9ac3abc10..9f82d9d7b 100644 --- a/spinifex/gateway/bedrock/weights.go +++ b/spinifex/gateway/bedrock/weights.go @@ -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 } diff --git a/spinifex/gateway/bedrock/weights_test.go b/spinifex/gateway/bedrock/weights_test.go index a790d405e..3317db98d 100644 --- a/spinifex/gateway/bedrock/weights_test.go +++ b/spinifex/gateway/bedrock/weights_test.go @@ -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) +}