mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 17:11:30 +00:00
feat: add spx admin ochre weights stage/list/remove (#680)
* feat: add spx admin ochre weights stage/list/remove Give operators a supported path to stage a self-hosted model's weights: point at a predastore S3 prefix holding ordinary Hugging Face files, and stage materialises them into a viperblock volume, snapshots it, and records the source URI and snapshot ID in the bedrock-weights KV bucket. stage refuses before downloading anything if the model is not a self-host catalog entry, or if the prefix is missing any required file. Re-staging the same source is a no-op; a different source replaces the entry and reports the previous snapshot for separate reclamation. WeightsStore gains GetWeights and a source-URI field on its record to support this. LookupServingSpec (catalog.go) is reused unchanged to validate --model-id. * fix: unify ServingSpec after rebase onto dev's daemon-launcher spec Rebasing onto dev pulled in mulga-vmfng.7.5's daemon-side ServingSpec and LookupServingSpec, textually duplicating the ochre weights stage CLI's own versions in catalog.go. Collapse both into one type carrying ModelID and VLLMArgs, and one three-value LookupServingSpec so callers distinguish an unknown model (found=false) from a provider-tier one (selfHost=false). Update the daemon launcher's call site accordingly. * test: cover ochre weights stage/list/remove command behaviour Splits runOchreWeightsStage/List/Remove into testable decision functions (runStageWeights, listWeightsOutput, removeWeights) separate from cobra and NATS wiring, and makes mkfs.ext4 execution and the connect/exit steps swappable so the full command flow can be driven against an embedded JetStream server instead of live infrastructure. Also relaxes the tokenizer file requirement to accept either tokenizer.json or tokenizer.model, since canonical Llama 3.x repos ship only the former at the prefix root. * fix: resolve mkfs.ext4 from sbin when it is absent from PATH e2fsprogs installs mkfs.ext4 into sbin, which a non-login shell omits, so weights staging failed on an unprivileged operator's otherwise fine host. * fix: size the weights volume to a whole GiB so the tail block fits viperblock rejects a write whose block overruns VolumeSize, so importing an image of unaligned length failed on its last block. Rounding up also makes the volume's byte size and its manifest SizeGiB describe the same volume. Name the volume in the import error so a part-written import is traceable.
This commit is contained in:
Josh Siviour
@@ -396,6 +396,8 @@ func init() {
|
||||
}
|
||||
}
|
||||
|
||||
const bytesPerGiB = 1024 * 1024 * 1024
|
||||
|
||||
// amiVolumeSizeGiB returns the smallest whole GiB that still holds sizeBytes.
|
||||
//
|
||||
// Rounding up is load-bearing. The image is copied into a root volume of
|
||||
@@ -410,7 +412,6 @@ func init() {
|
||||
// raises a caller's requested size to this value, so it inherits the mistake
|
||||
// rather than catching it.
|
||||
func amiVolumeSizeGiB(sizeBytes int64) uint64 {
|
||||
const bytesPerGiB = 1024 * 1024 * 1024
|
||||
if sizeBytes <= 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
@@ -0,0 +1,672 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/s3"
|
||||
"github.com/mulgadc/predastore/pkg/masterkey"
|
||||
"github.com/mulgadc/spinifex/spinifex/admin"
|
||||
"github.com/mulgadc/spinifex/spinifex/config"
|
||||
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
|
||||
"github.com/mulgadc/spinifex/spinifex/objectstore"
|
||||
"github.com/mulgadc/spinifex/spinifex/utils"
|
||||
"github.com/mulgadc/viperblock/viperblock"
|
||||
vbs3 "github.com/mulgadc/viperblock/viperblock/backends/s3"
|
||||
"github.com/mulgadc/viperblock/viperblock/v_utils"
|
||||
"github.com/nats-io/nats.go/jetstream"
|
||||
"github.com/pterm/pterm"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
var ochreCmd = &cobra.Command{
|
||||
Use: "ochre",
|
||||
Short: "Manage Ochre (self-hosted model inference) resources",
|
||||
}
|
||||
|
||||
var ochreWeightsCmd = &cobra.Command{
|
||||
Use: "weights",
|
||||
Short: "Stage, list and remove self-host model weights",
|
||||
Long: `Give an operator a supported path to stage a self-hosted model's weights so
|
||||
Ochre can advertise and serve it. A model with no staged weights is hidden
|
||||
from ListFoundationModels rather than advertised and broken.`,
|
||||
}
|
||||
|
||||
var ochreWeightsStageCmd = &cobra.Command{
|
||||
Use: "stage",
|
||||
Short: "Materialise a self-host model's weights from predastore into a servable snapshot",
|
||||
Long: `stage takes an S3 URI pointing at a Hugging Face model directory already
|
||||
uploaded to predastore (e.g. via 'aws s3 cp --recursive'), verifies the
|
||||
required files are present, materialises them into a viperblock volume,
|
||||
snapshots it, and records the source URI and snapshot ID against --model-id
|
||||
in the bedrock-weights KV bucket.
|
||||
|
||||
Idempotent: re-staging the same --s3-uri for a model that already has it
|
||||
staged is a no-op. Re-staging a different --s3-uri replaces the KV entry and
|
||||
reports the previous snapshot ID so an operator can reclaim it separately.
|
||||
|
||||
Refuses before materialising anything if --model-id is not a self-host
|
||||
catalog entry, or if the S3 prefix is missing any required file.`,
|
||||
Run: runOchreWeightsStage,
|
||||
}
|
||||
|
||||
var ochreWeightsListCmd = &cobra.Command{
|
||||
Use: "list",
|
||||
Short: "List staged self-host model weights",
|
||||
Run: runOchreWeightsList,
|
||||
}
|
||||
|
||||
var ochreWeightsRemoveCmd = &cobra.Command{
|
||||
Use: "remove",
|
||||
Short: "Drop a model's staged-weights KV entry",
|
||||
Long: `remove drops --model-id's entry from the bedrock-weights KV bucket, which
|
||||
hides it from ListFoundationModels again. It never deletes the underlying
|
||||
snapshot or the source S3 objects; reclaiming that storage is a separate,
|
||||
explicit act.`,
|
||||
Run: runOchreWeightsRemove,
|
||||
}
|
||||
|
||||
func init() {
|
||||
adminCmd.AddCommand(ochreCmd)
|
||||
ochreCmd.AddCommand(ochreWeightsCmd)
|
||||
ochreWeightsCmd.AddCommand(ochreWeightsStageCmd)
|
||||
ochreWeightsCmd.AddCommand(ochreWeightsListCmd)
|
||||
ochreWeightsCmd.AddCommand(ochreWeightsRemoveCmd)
|
||||
|
||||
ochreWeightsStageCmd.Flags().String("model-id", "", "Catalog model ID to stage weights for (required)")
|
||||
ochreWeightsStageCmd.Flags().String("s3-uri", "", "predastore S3 URI holding the model's Hugging Face files, e.g. s3://bucket/prefix/ (required)")
|
||||
ochreWeightsStageCmd.Flags().String("tmp-dir", os.TempDir(), "Temporary directory for download and volume staging")
|
||||
_ = ochreWeightsStageCmd.MarkFlagRequired("model-id")
|
||||
_ = ochreWeightsStageCmd.MarkFlagRequired("s3-uri")
|
||||
|
||||
ochreWeightsRemoveCmd.Flags().String("model-id", "", "Model ID to remove from staged weights (required)")
|
||||
_ = ochreWeightsRemoveCmd.MarkFlagRequired("model-id")
|
||||
}
|
||||
|
||||
// requiredWeightsFiles are the fixed-name Hugging Face artefacts stage
|
||||
// refuses to materialise without, mirroring what AWS Bedrock's
|
||||
// CreateModelImportJob expects at its S3 source prefix. At least one
|
||||
// *.safetensors file and one tokenizer file are checked separately.
|
||||
var requiredWeightsFiles = []string{
|
||||
"config.json",
|
||||
"tokenizer_config.json",
|
||||
}
|
||||
|
||||
// tokenizerFileNames are the two shapes a Hugging Face repo ships its
|
||||
// tokenizer under: modern repos ship only tokenizer.json (fast tokenizer),
|
||||
// older ones only tokenizer.model (SentencePiece) -- e.g. Llama 3.x keeps
|
||||
// tokenizer.model under original/, not at the prefix root. Either is enough.
|
||||
var tokenizerFileNames = []string{"tokenizer.json", "tokenizer.model"}
|
||||
|
||||
// parseWeightsS3URI splits an s3://bucket/prefix URI into its bucket and
|
||||
// prefix. The prefix is normalised to end with '/' so downstream listing and
|
||||
// validation scope to the directory rather than any key merely sharing the
|
||||
// prefix as a substring.
|
||||
func parseWeightsS3URI(uri string) (bucket, prefix string, err error) {
|
||||
trimmed := strings.TrimPrefix(uri, "s3://")
|
||||
if trimmed == uri {
|
||||
return "", "", fmt.Errorf("invalid --s3-uri %q: expected s3://bucket/prefix", uri)
|
||||
}
|
||||
parts := strings.SplitN(trimmed, "/", 2)
|
||||
bucket = parts[0]
|
||||
if bucket == "" {
|
||||
return "", "", fmt.Errorf("invalid --s3-uri %q: missing bucket", uri)
|
||||
}
|
||||
if len(parts) == 2 {
|
||||
prefix = parts[1]
|
||||
}
|
||||
if prefix != "" && !strings.HasSuffix(prefix, "/") {
|
||||
prefix += "/"
|
||||
}
|
||||
return bucket, prefix, nil
|
||||
}
|
||||
|
||||
// listWeightsPrefix pages through every object under bucket/prefix, calling
|
||||
// fn for each. Both validateWeightsPrefix (presence-only) and
|
||||
// downloadWeightsPrefix (full copy) share this walk.
|
||||
func listWeightsPrefix(ctx context.Context, store objectstore.ObjectStore, bucket, prefix string, fn func(*s3.Object) error) error {
|
||||
var token *string
|
||||
for {
|
||||
out, err := store.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(bucket),
|
||||
Prefix: aws.String(prefix),
|
||||
ContinuationToken: token,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, obj := range out.Contents {
|
||||
if err := fn(obj); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if out.IsTruncated == nil || !*out.IsTruncated || out.NextContinuationToken == nil {
|
||||
return nil
|
||||
}
|
||||
token = out.NextContinuationToken
|
||||
}
|
||||
}
|
||||
|
||||
// validateWeightsPrefix confirms every required Hugging Face file exists
|
||||
// under bucket/prefix before stage downloads anything, so a typo'd prefix
|
||||
// fails in milliseconds rather than after materialising multiple gigabytes.
|
||||
func validateWeightsPrefix(ctx context.Context, store objectstore.ObjectStore, bucket, prefix string) error {
|
||||
var missing []string
|
||||
for _, name := range requiredWeightsFiles {
|
||||
key := prefix + name
|
||||
if _, err := store.HeadObject(ctx, &s3.HeadObjectInput{Bucket: aws.String(bucket), Key: aws.String(key)}); err != nil {
|
||||
if objectstore.IsNoSuchKeyError(err) {
|
||||
missing = append(missing, name)
|
||||
continue
|
||||
}
|
||||
return fmt.Errorf("head s3://%s/%s: %w", bucket, key, err)
|
||||
}
|
||||
}
|
||||
|
||||
hasSafetensors := false
|
||||
if err := listWeightsPrefix(ctx, store, bucket, prefix, func(obj *s3.Object) error {
|
||||
if strings.HasSuffix(aws.StringValue(obj.Key), ".safetensors") {
|
||||
hasSafetensors = true
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return fmt.Errorf("list s3://%s/%s: %w", bucket, prefix, err)
|
||||
}
|
||||
if !hasSafetensors {
|
||||
missing = append(missing, "*.safetensors")
|
||||
}
|
||||
|
||||
hasTokenizer := false
|
||||
for _, name := range tokenizerFileNames {
|
||||
key := prefix + name
|
||||
if _, err := store.HeadObject(ctx, &s3.HeadObjectInput{Bucket: aws.String(bucket), Key: aws.String(key)}); err == nil {
|
||||
hasTokenizer = true
|
||||
break
|
||||
} else if !objectstore.IsNoSuchKeyError(err) {
|
||||
return fmt.Errorf("head s3://%s/%s: %w", bucket, key, err)
|
||||
}
|
||||
}
|
||||
if !hasTokenizer {
|
||||
missing = append(missing, "tokenizer.json or tokenizer.model")
|
||||
}
|
||||
|
||||
if len(missing) > 0 {
|
||||
return fmt.Errorf("s3://%s/%s is missing required file(s): %s", bucket, prefix, strings.Join(missing, ", "))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// downloadWeightsPrefix copies every object under bucket/prefix into
|
||||
// destDir, flattening predastore's key structure to the basename: a Hugging
|
||||
// Face model directory is flat, and the filesystem image stage builds only
|
||||
// needs the files, not the source key layout.
|
||||
func downloadWeightsPrefix(ctx context.Context, store objectstore.ObjectStore, bucket, prefix, destDir string) (int64, error) {
|
||||
var total int64
|
||||
err := listWeightsPrefix(ctx, store, bucket, prefix, func(obj *s3.Object) error {
|
||||
key := aws.StringValue(obj.Key)
|
||||
if strings.HasSuffix(key, "/") {
|
||||
return nil // directory marker, not a file
|
||||
}
|
||||
n, err := downloadObjectTo(ctx, store, bucket, key, filepath.Join(destDir, path.Base(key)))
|
||||
if err != nil {
|
||||
return fmt.Errorf("download s3://%s/%s: %w", bucket, key, err)
|
||||
}
|
||||
total += n
|
||||
return nil
|
||||
})
|
||||
return total, err
|
||||
}
|
||||
|
||||
func downloadObjectTo(ctx context.Context, store objectstore.ObjectStore, bucket, key, destPath string) (int64, error) {
|
||||
out, err := store.GetObject(ctx, &s3.GetObjectInput{Bucket: aws.String(bucket), Key: aws.String(key)})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer out.Body.Close()
|
||||
|
||||
f, err := os.OpenFile(destPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
return io.Copy(f, out.Body)
|
||||
}
|
||||
|
||||
// mkfsExt4SbinDirs are searched when mkfs.ext4 is absent from PATH. e2fsprogs
|
||||
// installs into sbin, which a non-login shell routinely omits, so a plain
|
||||
// LookPath fails for an unprivileged operator on an otherwise fine host.
|
||||
var mkfsExt4SbinDirs = []string{"/usr/local/sbin", "/usr/sbin", "/sbin"}
|
||||
|
||||
// lookupMkfsExt4 resolves mkfs.ext4 from PATH, falling back to the sbin dirs.
|
||||
func lookupMkfsExt4() (string, error) {
|
||||
if path, err := exec.LookPath("mkfs.ext4"); err == nil {
|
||||
return path, nil
|
||||
}
|
||||
for _, dir := range mkfsExt4SbinDirs {
|
||||
candidate := filepath.Join(dir, "mkfs.ext4")
|
||||
if info, err := os.Stat(candidate); err == nil && !info.IsDir() {
|
||||
return candidate, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("mkfs.ext4 not found on PATH or in %s (install e2fsprogs)",
|
||||
strings.Join(mkfsExt4SbinDirs, ", "))
|
||||
}
|
||||
|
||||
// mkfsExt4Runner populates imagePath as an ext4 filesystem from srcDir's
|
||||
// files, normally by shelling out to mkfs.ext4 -d. A package var, mirroring
|
||||
// caBakeRunner, so tests can substitute a fake instead of requiring
|
||||
// mkfs.ext4 on PATH.
|
||||
var mkfsExt4Runner = func(srcDir, imagePath string) error {
|
||||
mkfs, err := lookupMkfsExt4()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
out, err := exec.Command(mkfs, "-F", "-d", srcDir, imagePath).CombinedOutput()
|
||||
if err != nil {
|
||||
return fmt.Errorf("mkfs.ext4: %w: %s", err, string(out))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// buildWeightsImage packages srcDir's files into a raw ext4 filesystem image
|
||||
// at imagePath, sized to fit their total bytes plus filesystem overhead and
|
||||
// headroom. mkfs.ext4 -d populates the filesystem directly from srcDir, so
|
||||
// no loopback mount (and no root) is needed -- unlike the guestfish/
|
||||
// virt-customize tooling build-system-image.sh needs to customize a
|
||||
// bootable cloud image, a weights volume is just a directory of files.
|
||||
func buildWeightsImage(srcDir, imagePath string, contentBytes int64) error {
|
||||
const overheadFraction = 0.15 // ext4 metadata + inode table headroom
|
||||
const minPaddingBytes = 64 * 1024 * 1024
|
||||
padding := max(int64(float64(contentBytes)*overheadFraction), minPaddingBytes)
|
||||
sizeBytes := contentBytes + padding
|
||||
|
||||
f, err := os.OpenFile(imagePath, os.O_CREATE|os.O_WRONLY, 0600)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create image file: %w", err)
|
||||
}
|
||||
if err := f.Truncate(sizeBytes); err != nil {
|
||||
f.Close()
|
||||
return fmt.Errorf("size image file: %w", err)
|
||||
}
|
||||
if err := f.Close(); err != nil {
|
||||
return fmt.Errorf("close image file: %w", err)
|
||||
}
|
||||
|
||||
return mkfsExt4Runner(srcDir, imagePath)
|
||||
}
|
||||
|
||||
// snapshotImportedWeightsVolume snapshots a viperblock volume that
|
||||
// v_utils.ImportDiskImage just wrote and closed -- never attached, never
|
||||
// nbdkit-served, so there is no live checkpoint to load. This mirrors the
|
||||
// offline snapshot sequence handlers/ec2/image/service_impl.go's
|
||||
// snapshotStoppedVolume uses for a stopped instance's root volume: reopen
|
||||
// read-only, load the numbered checkpoint Close() wrote, then CreateSnapshot.
|
||||
func snapshotImportedWeightsVolume(s3Config *vbs3.S3Config, volumeID string, volumeSize uint64, walDir string, mkey *masterkey.Key) (string, error) {
|
||||
vbConfig := viperblock.VB{
|
||||
VolumeName: volumeID,
|
||||
VolumeSize: volumeSize,
|
||||
BaseDir: walDir,
|
||||
Cache: viperblock.Cache{Config: viperblock.CacheConfig{Size: 0}},
|
||||
MasterKey: mkey,
|
||||
EncryptionEnabled: mkey != nil,
|
||||
Logger: slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})),
|
||||
}
|
||||
|
||||
vb, err := viperblock.New(&vbConfig, "s3", *s3Config)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("new viperblock: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
vb.StopChunkUploader()
|
||||
vb.StopWALSyncer()
|
||||
}()
|
||||
|
||||
if err := vb.Backend.Init(); err != nil {
|
||||
return "", fmt.Errorf("backend init: %w", err)
|
||||
}
|
||||
if err := vb.LoadState(); err != nil {
|
||||
return "", fmt.Errorf("load state: %w", err)
|
||||
}
|
||||
if err := vb.LoadBlockState(); err != nil {
|
||||
return "", fmt.Errorf("load block state: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := vb.RemoveLocalFiles(); err != nil {
|
||||
slog.Warn("snapshotImportedWeightsVolume: failed to remove local files", "volumeId", volumeID, "err", err)
|
||||
}
|
||||
}()
|
||||
|
||||
snapshotID := admin.SnapPrefix(volumeID)
|
||||
if _, err := vb.CreateSnapshot(snapshotID); err != nil {
|
||||
return "", fmt.Errorf("create snapshot: %w", err)
|
||||
}
|
||||
return snapshotID, nil
|
||||
}
|
||||
|
||||
// weightsMaterializer builds a servable snapshot from downloadDir's already
|
||||
// downloaded and validated Hugging Face files. Isolating this behind a
|
||||
// function value lets runStageWeights be tested end to end -- including the
|
||||
// idempotent-noop and replace-and-report-previous-snapshot decisions -- with
|
||||
// a fake standing in for the real mkfs.ext4 + viperblock + predastore work,
|
||||
// which needs a live environment (see materializeWeightsVolume).
|
||||
type weightsMaterializer func(ctx context.Context, downloadDir string, contentBytes int64) (snapshotID string, err error)
|
||||
|
||||
// runStageWeights holds 'ochre weights stage' decision and side-effect logic:
|
||||
// catalog validation, S3 URI parsing, idempotency/replacement detection
|
||||
// against the existing KV record, prefix validation, download, materialisation
|
||||
// via materialize, and the final KV write. It has no cobra or NATS-connection
|
||||
// dependency of its own, so tests drive it directly against a fake object
|
||||
// store, a real WeightsStore backed by an embedded JetStream server, and a
|
||||
// fake materializer.
|
||||
func runStageWeights(ctx context.Context, store objectstore.ObjectStore, weightsStore *gateway_bedrock.WeightsStore, tmpDirFlag, modelID, s3URI string, materialize weightsMaterializer) (string, error) {
|
||||
if _, found, selfHost := gateway_bedrock.LookupServingSpec(modelID); !found || !selfHost {
|
||||
if !found {
|
||||
return "", fmt.Errorf("unknown model ID %q: not present in the Ochre catalog", modelID)
|
||||
}
|
||||
return "", fmt.Errorf("%q is a provider-served model, not self-host; weights staging does not apply", modelID)
|
||||
}
|
||||
|
||||
bucket, prefix, err := parseWeightsS3URI(s3URI)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
sourceURI := fmt.Sprintf("s3://%s/%s", bucket, prefix)
|
||||
|
||||
existing, hadPrevious, err := weightsStore.GetWeights(ctx, modelID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if hadPrevious && existing.SourceURI == sourceURI {
|
||||
return fmt.Sprintf("%s is already staged from %s (snapshot %s); nothing to do.", modelID, sourceURI, existing.SnapshotID), nil
|
||||
}
|
||||
|
||||
fmt.Printf("Validating %s ...\n", sourceURI)
|
||||
if err := validateWeightsPrefix(ctx, store, bucket, prefix); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
tmpDir, err := os.MkdirTemp(tmpDirFlag, "spinifex-weights-tmp-*")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("could not create temp dir: %w", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
downloadDir := filepath.Join(tmpDir, "src")
|
||||
if err := os.MkdirAll(downloadDir, 0700); err != nil {
|
||||
return "", fmt.Errorf("could not create download dir: %w", err)
|
||||
}
|
||||
|
||||
fmt.Printf("Downloading %s ...\n", sourceURI)
|
||||
contentBytes, err := downloadWeightsPrefix(ctx, store, bucket, prefix, downloadDir)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
snapshotID, err := materialize(ctx, downloadDir, contentBytes)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if err := weightsStore.PutWeights(ctx, modelID, sourceURI, snapshotID); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if hadPrevious {
|
||||
return fmt.Sprintf("✅ Staged %s from %s (snapshot %s). Replaced previous snapshot %s -- reclaim it separately if no longer needed.",
|
||||
modelID, sourceURI, snapshotID, existing.SnapshotID), nil
|
||||
}
|
||||
return fmt.Sprintf("✅ Staged %s from %s (snapshot %s).", modelID, sourceURI, snapshotID), nil
|
||||
}
|
||||
|
||||
// materializeWeightsVolume is the real weightsMaterializer: it packages
|
||||
// downloadDir into a filesystem image, imports it into a new viperblock
|
||||
// volume backed by node's predastore, and snapshots it. This is the
|
||||
// expensive, environment-dependent half of stage -- it shells out to
|
||||
// mkfs.ext4 and talks to a live S3-shaped backend -- so it is exercised live
|
||||
// (see the plan's Testing section) rather than in unit tests.
|
||||
func materializeWeightsVolume(node config.Config, downloadDir string, contentBytes int64, mkey *masterkey.Key) (string, error) {
|
||||
tmpDir := filepath.Dir(downloadDir)
|
||||
volumeId := utils.GenerateResourceID("vol")
|
||||
imagePath := filepath.Join(tmpDir, volumeId+".img")
|
||||
|
||||
fmt.Println("Building filesystem image ...")
|
||||
if err := buildWeightsImage(downloadDir, imagePath, contentBytes); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
imageStat, err := os.Stat(imagePath)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("could not stat image: %w", err)
|
||||
}
|
||||
|
||||
// Round the volume up to a whole GiB rather than using the raw image size.
|
||||
// viperblock rejects a tail write whose block overruns VolumeSize, so an
|
||||
// unaligned size fails at the last block. This also keeps the byte size and
|
||||
// the manifest's SizeGiB describing the same volume.
|
||||
sizeGiB := amiVolumeSizeGiB(imageStat.Size())
|
||||
volumeBytes := sizeGiB * bytesPerGiB
|
||||
|
||||
s3Config := vbs3.S3Config{
|
||||
VolumeName: volumeId,
|
||||
VolumeSize: volumeBytes,
|
||||
Bucket: node.Predastore.Bucket,
|
||||
Region: node.Predastore.Region,
|
||||
AccessKey: node.Predastore.AccessKey,
|
||||
SecretKey: node.Predastore.SecretKey,
|
||||
Host: node.Predastore.Host,
|
||||
}
|
||||
|
||||
// AMIMetadata is deliberately left zero-valued: a weights volume is not
|
||||
// bootable and must never be registered as a launchable AMI. Leaving it
|
||||
// unset also means ImportDiskImage does not attempt its own automatic
|
||||
// snapshot -- that happens explicitly below, after the volume is closed.
|
||||
manifest := viperblock.VolumeConfig{}
|
||||
manifest.VolumeMetadata.VolumeID = volumeId
|
||||
manifest.VolumeMetadata.VolumeName = volumeId
|
||||
manifest.VolumeMetadata.TenantID = "system"
|
||||
manifest.VolumeMetadata.SizeGiB = sizeGiB
|
||||
manifest.VolumeMetadata.State = "available"
|
||||
manifest.VolumeMetadata.AvailabilityZone = node.Predastore.Region
|
||||
manifest.VolumeMetadata.CreatedAt = time.Now()
|
||||
manifest.VolumeMetadata.VolumeType = "gp3"
|
||||
manifest.VolumeMetadata.IOPS = 1000
|
||||
|
||||
vbConfig := viperblock.VB{
|
||||
VolumeName: volumeId,
|
||||
VolumeSize: volumeBytes,
|
||||
BaseDir: tmpDir,
|
||||
Cache: viperblock.Cache{Config: viperblock.CacheConfig{Size: 0}},
|
||||
VolumeConfig: manifest,
|
||||
MasterKey: mkey,
|
||||
EncryptionEnabled: mkey != nil,
|
||||
Logger: slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError})),
|
||||
}
|
||||
|
||||
var flushBar *pterm.ProgressbarPrinter
|
||||
var flushUpdate func(current uint64)
|
||||
progress := func(current, total uint64) {
|
||||
if flushBar == nil {
|
||||
flushBar, flushUpdate = utils.NewByteProgressBar("Flushing weights to storage", total)
|
||||
}
|
||||
flushUpdate(current)
|
||||
}
|
||||
|
||||
// Name the volume in the failure path: a part-written import leaves blocks
|
||||
// in predastore, and without the ID an operator cannot find them to reclaim.
|
||||
if err := v_utils.ImportDiskImage(&s3Config, &vbConfig, imagePath, progress); err != nil {
|
||||
if flushBar != nil {
|
||||
_, _ = flushBar.Stop()
|
||||
}
|
||||
return "", fmt.Errorf("could not import weights volume %s: %w", volumeId, err)
|
||||
}
|
||||
if flushBar != nil {
|
||||
_, _ = flushBar.Stop()
|
||||
}
|
||||
|
||||
fmt.Println("Snapshotting weights volume ...")
|
||||
snapshotID, err := snapshotImportedWeightsVolume(&s3Config, volumeId, volumeBytes, tmpDir, mkey)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("could not snapshot weights volume: %w", err)
|
||||
}
|
||||
return snapshotID, nil
|
||||
}
|
||||
|
||||
// loadConfigAndConnectFn and ochreExit indirect the two effects that make
|
||||
// the ochre weights Run functions otherwise untestable: a live NATS
|
||||
// connection and a process-terminating exit. Tests substitute both -- a fake
|
||||
// connection backed by an embedded JetStream server, and an exit stand-in
|
||||
// that panics with a sentinel instead of killing the test binary -- so the
|
||||
// same connect/validate/exit control flow real operators see is exercised
|
||||
// directly.
|
||||
var (
|
||||
loadConfigAndConnectFn = loadConfigAndConnect
|
||||
ochreExit = os.Exit
|
||||
)
|
||||
|
||||
func runOchreWeightsStage(cmd *cobra.Command, _ []string) {
|
||||
modelID, _ := cmd.Flags().GetString("model-id")
|
||||
s3URI, _ := cmd.Flags().GetString("s3-uri")
|
||||
tmpDirFlag, _ := cmd.Flags().GetString("tmp-dir")
|
||||
|
||||
appConfig, nc, err := loadConfigAndConnectFn()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
defer nc.Close()
|
||||
node := appConfig.Nodes[appConfig.Node]
|
||||
|
||||
js, err := jetstream.New(nc)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
weightsStore := gateway_bedrock.NewWeightsStore(js, len(appConfig.Nodes))
|
||||
store := objectstore.NewS3ObjectStoreFromConfig(node.Predastore.Host, node.Predastore.Region, node.Predastore.AccessKey, node.Predastore.SecretKey)
|
||||
|
||||
mkey, err := utils.LoadViperblockMasterKey(node.Viperblock.EncryptionKeyFile)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Could not load viperblock encryption key: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
|
||||
materialize := func(_ context.Context, downloadDir string, contentBytes int64) (string, error) {
|
||||
return materializeWeightsVolume(node, downloadDir, contentBytes, mkey)
|
||||
}
|
||||
|
||||
msg, err := runStageWeights(context.Background(), store, weightsStore, tmpDirFlag, modelID, s3URI, materialize)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "%v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
fmt.Println(msg)
|
||||
}
|
||||
|
||||
// listWeightsOutput renders 'ochre weights list': a friendly message when
|
||||
// nothing is staged, else a MODEL ID / SOURCE URI / SNAPSHOT ID table. Split
|
||||
// out from runOchreWeightsList so it is testable against a real WeightsStore
|
||||
// without a NATS connection.
|
||||
func listWeightsOutput(ctx context.Context, weightsStore *gateway_bedrock.WeightsStore) (string, error) {
|
||||
entries, err := weightsStore.ListWeights(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(entries) == 0 {
|
||||
return "No models staged.", nil
|
||||
}
|
||||
|
||||
tableData := pterm.TableData{{"MODEL ID", "SOURCE URI", "SNAPSHOT ID"}}
|
||||
for _, e := range entries {
|
||||
tableData = append(tableData, []string{e.ModelID, e.SourceURI, e.SnapshotID})
|
||||
}
|
||||
return pterm.DefaultTable.WithHasHeader().WithLeftAlignment().WithData(tableData).Srender()
|
||||
}
|
||||
|
||||
func runOchreWeightsList(_ *cobra.Command, _ []string) {
|
||||
appConfig, nc, err := loadConfigAndConnectFn()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
defer nc.Close()
|
||||
|
||||
js, err := jetstream.New(nc)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
weightsStore := gateway_bedrock.NewWeightsStore(js, len(appConfig.Nodes))
|
||||
|
||||
out, err := listWeightsOutput(context.Background(), weightsStore)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
fmt.Println(out)
|
||||
}
|
||||
|
||||
// removeWeights drops modelID's KV entry after confirming it exists, so the
|
||||
// CLI can report a clear "nothing staged" error rather than a generic KV
|
||||
// miss. It never touches the backing snapshot or source S3 objects -- only
|
||||
// weightsStore.DeleteWeights's KV key is affected. Split out from
|
||||
// runOchreWeightsRemove so it is testable against a real WeightsStore
|
||||
// without a NATS connection.
|
||||
func removeWeights(ctx context.Context, weightsStore *gateway_bedrock.WeightsStore, modelID string) (gateway_bedrock.WeightsEntry, error) {
|
||||
entry, ok, err := weightsStore.GetWeights(ctx, modelID)
|
||||
if err != nil {
|
||||
return gateway_bedrock.WeightsEntry{}, err
|
||||
}
|
||||
if !ok {
|
||||
return gateway_bedrock.WeightsEntry{}, fmt.Errorf("%s has no staged weights entry", modelID)
|
||||
}
|
||||
|
||||
if err := weightsStore.DeleteWeights(ctx, modelID); err != nil {
|
||||
return gateway_bedrock.WeightsEntry{}, err
|
||||
}
|
||||
return entry, nil
|
||||
}
|
||||
|
||||
func runOchreWeightsRemove(cmd *cobra.Command, _ []string) {
|
||||
modelID, _ := cmd.Flags().GetString("model-id")
|
||||
|
||||
appConfig, nc, err := loadConfigAndConnectFn()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
defer nc.Close()
|
||||
|
||||
js, err := jetstream.New(nc)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
weightsStore := gateway_bedrock.NewWeightsStore(js, len(appConfig.Nodes))
|
||||
|
||||
entry, err := removeWeights(context.Background(), weightsStore, modelID)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "%v\n", err)
|
||||
ochreExit(1)
|
||||
return
|
||||
}
|
||||
fmt.Printf("✅ Removed staged-weights entry for %s (snapshot %s and source objects untouched).\n", modelID, entry.SnapshotID)
|
||||
}
|
||||
@@ -0,0 +1,951 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/aws/aws-sdk-go/aws"
|
||||
"github.com/aws/aws-sdk-go/service/s3"
|
||||
"github.com/mulgadc/spinifex/spinifex/config"
|
||||
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
|
||||
"github.com/mulgadc/spinifex/spinifex/objectstore"
|
||||
"github.com/mulgadc/spinifex/spinifex/testutil"
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// selfHostModelID and providerModelID are real catalog entries: the former
|
||||
// is self-host (weights staging applies), the latter is provider-served
|
||||
// (weights staging must refuse it), so tests exercise LookupServingSpec's
|
||||
// found/selfHost distinction against the real catalog rather than a stub.
|
||||
const (
|
||||
selfHostModelID = "meta.llama3-2-1b-instruct-v1:0"
|
||||
providerModelID = "anthropic.claude-3-5-sonnet-20240620-v1:0"
|
||||
)
|
||||
|
||||
// newWeightsStoreForTest builds a WeightsStore backed by an embedded
|
||||
// JetStream server, so stage/list/remove tests exercise the real KV
|
||||
// idempotency and persistence logic instead of a hand-rolled fake.
|
||||
func newWeightsStoreForTest(t *testing.T) *gateway_bedrock.WeightsStore {
|
||||
t.Helper()
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
return gateway_bedrock.NewWeightsStore(testutil.NewJetStream(t, nc), 1)
|
||||
}
|
||||
|
||||
// explodingObjectStore fails the test immediately if any of its methods is
|
||||
// called. It stands in for the object store on paths that must do zero S3
|
||||
// work, such as the idempotent no-op re-stage and the two catalog-refusal
|
||||
// cases.
|
||||
type explodingObjectStore struct {
|
||||
t *testing.T
|
||||
}
|
||||
|
||||
func (e explodingObjectStore) GetObject(context.Context, *s3.GetObjectInput) (*s3.GetObjectOutput, error) {
|
||||
e.t.Fatal("GetObject called on a path that must not touch S3")
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (e explodingObjectStore) HeadObject(context.Context, *s3.HeadObjectInput) (*s3.HeadObjectOutput, error) {
|
||||
e.t.Fatal("HeadObject called on a path that must not touch S3")
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (e explodingObjectStore) PutObject(context.Context, *s3.PutObjectInput) (*s3.PutObjectOutput, error) {
|
||||
e.t.Fatal("PutObject called on a path that must not touch S3")
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (e explodingObjectStore) DeleteObject(context.Context, *s3.DeleteObjectInput) (*s3.DeleteObjectOutput, error) {
|
||||
e.t.Fatal("DeleteObject called on a path that must not touch S3")
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (e explodingObjectStore) ListObjectsV2(context.Context, *s3.ListObjectsV2Input) (*s3.ListObjectsV2Output, error) {
|
||||
e.t.Fatal("ListObjectsV2 called on a path that must not touch S3")
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (e explodingObjectStore) EnsureBucket(context.Context, string) error {
|
||||
e.t.Fatal("EnsureBucket called on a path that must not touch S3")
|
||||
return nil
|
||||
}
|
||||
|
||||
var _ objectstore.ObjectStore = explodingObjectStore{}
|
||||
|
||||
// explodingMaterializer fails the test immediately if invoked, standing in
|
||||
// for weightsMaterializer on paths that must refuse before materialising
|
||||
// anything.
|
||||
func explodingMaterializer(t *testing.T) weightsMaterializer {
|
||||
t.Helper()
|
||||
return func(context.Context, string, int64) (string, error) {
|
||||
t.Fatal("materialize called on a path that must refuse before materialising")
|
||||
return "", nil
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseWeightsS3URI_ValidWithTrailingSlash(t *testing.T) {
|
||||
bucket, prefix, err := parseWeightsS3URI("s3://models/llama-3.2-1b/")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "models", bucket)
|
||||
assert.Equal(t, "llama-3.2-1b/", prefix)
|
||||
}
|
||||
|
||||
// TestParseWeightsS3URI_AddsMissingTrailingSlash covers scoping: without the
|
||||
// trailing slash, downstream prefix listing would also match a sibling
|
||||
// object whose key merely shares the prefix as a substring.
|
||||
func TestParseWeightsS3URI_AddsMissingTrailingSlash(t *testing.T) {
|
||||
_, prefix, err := parseWeightsS3URI("s3://models/llama-3.2-1b")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "llama-3.2-1b/", prefix)
|
||||
}
|
||||
|
||||
func TestParseWeightsS3URI_BucketOnlyNoPrefix(t *testing.T) {
|
||||
bucket, prefix, err := parseWeightsS3URI("s3://models")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "models", bucket)
|
||||
assert.Empty(t, prefix)
|
||||
}
|
||||
|
||||
func TestParseWeightsS3URI_MissingSchemeIsError(t *testing.T) {
|
||||
_, _, err := parseWeightsS3URI("models/llama-3.2-1b/")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestParseWeightsS3URI_MissingBucketIsError(t *testing.T) {
|
||||
_, _, err := parseWeightsS3URI("s3:///llama-3.2-1b/")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
// putObject seeds a memory object store with key -> body, as if it were an
|
||||
// already-uploaded Hugging Face model file.
|
||||
func putObject(t *testing.T, store *objectstore.MemoryObjectStore, bucket, key string, body []byte) {
|
||||
t.Helper()
|
||||
_, err := store.PutObject(context.Background(), &s3.PutObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(key),
|
||||
Body: bytes.NewReader(body),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func seedCompleteWeightsPrefix(t *testing.T, store *objectstore.MemoryObjectStore, bucket, prefix string) {
|
||||
t.Helper()
|
||||
putObject(t, store, bucket, prefix+"model.safetensors", nil)
|
||||
putObject(t, store, bucket, prefix+"tokenizer.json", nil)
|
||||
for _, name := range requiredWeightsFiles {
|
||||
putObject(t, store, bucket, prefix+name, nil)
|
||||
}
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_AllFilesPresent covers the happy path: every
|
||||
// required file and at least one *.safetensors file present passes.
|
||||
func TestValidateWeightsPrefix_AllFilesPresent(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
seedCompleteWeightsPrefix(t, store, "models", "llama-3.2-1b/")
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_MissingRequiredFile covers refusal before any
|
||||
// materialisation: a typo'd or incomplete prefix must fail validation, not
|
||||
// be discovered after downloading gigabytes.
|
||||
func TestValidateWeightsPrefix_MissingRequiredFile(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
putObject(t, store, "models", "llama-3.2-1b/model.safetensors", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/config.json", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/tokenizer.json", nil)
|
||||
// tokenizer_config.json deliberately omitted.
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "tokenizer_config.json")
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_TokenizerJSONOnlyPasses covers modern
|
||||
// fast-tokenizer repos such as Llama 3.x, which ship only tokenizer.json at
|
||||
// the prefix root (tokenizer.model, if present at all, lives under
|
||||
// original/ and is not required).
|
||||
func TestValidateWeightsPrefix_TokenizerJSONOnlyPasses(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
putObject(t, store, "models", "llama-3.2-1b/model.safetensors", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/config.json", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/tokenizer_config.json", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/tokenizer.json", nil)
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_TokenizerModelOnlyPasses covers older
|
||||
// SentencePiece repos, which ship only tokenizer.model with no
|
||||
// tokenizer.json at all.
|
||||
func TestValidateWeightsPrefix_TokenizerModelOnlyPasses(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
putObject(t, store, "models", "llama-3.2-1b/model.safetensors", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/config.json", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/tokenizer_config.json", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/tokenizer.model", nil)
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_NeitherTokenizerFilePresent covers refusal when
|
||||
// neither tokenizer shape is present: the error must name the requirement.
|
||||
func TestValidateWeightsPrefix_NeitherTokenizerFilePresent(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
putObject(t, store, "models", "llama-3.2-1b/model.safetensors", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/config.json", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/tokenizer_config.json", nil)
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "tokenizer.json")
|
||||
assert.ErrorContains(t, err, "tokenizer.model")
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_MissingSafetensors covers the variable-name
|
||||
// weights file: all fixed-name files present is not enough without at least
|
||||
// one *.safetensors object.
|
||||
func TestValidateWeightsPrefix_MissingSafetensors(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
putObject(t, store, "models", "llama-3.2-1b/tokenizer.json", nil)
|
||||
for _, name := range requiredWeightsFiles {
|
||||
putObject(t, store, "models", "llama-3.2-1b/"+name, nil)
|
||||
}
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "*.safetensors")
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_EmptyPrefixIsAllMissing covers a wrong/typo'd
|
||||
// prefix that resolves to nothing: every required file (and *.safetensors)
|
||||
// is reported missing, rather than a bare "not found".
|
||||
func TestValidateWeightsPrefix_EmptyPrefixIsAllMissing(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "does-not-exist/")
|
||||
require.Error(t, err)
|
||||
for _, name := range requiredWeightsFiles {
|
||||
assert.ErrorContains(t, err, name)
|
||||
}
|
||||
assert.ErrorContains(t, err, "*.safetensors")
|
||||
assert.ErrorContains(t, err, "tokenizer.json")
|
||||
assert.ErrorContains(t, err, "tokenizer.model")
|
||||
}
|
||||
|
||||
// TestDownloadWeightsPrefix_FlattensToBasename covers the download step:
|
||||
// predastore's key structure is flattened to the file's basename in destDir,
|
||||
// since buildWeightsImage only needs a flat directory of files.
|
||||
func TestDownloadWeightsPrefix_FlattensToBasename(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
content := []byte("weights-bytes")
|
||||
putObject(t, store, "models", "llama-3.2-1b/model.safetensors", content)
|
||||
putObject(t, store, "models", "llama-3.2-1b/config.json", []byte("{}"))
|
||||
|
||||
destDir := t.TempDir()
|
||||
total, err := downloadWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/", destDir)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(len(content)+len("{}")), total)
|
||||
|
||||
got, err := os.ReadFile(filepath.Join(destDir, "model.safetensors"))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, content, got)
|
||||
}
|
||||
|
||||
// TestRunStageWeights_IdempotentReStageIsNoop covers the no-op path: staging
|
||||
// the same --s3-uri a second time must not touch the object store or
|
||||
// materialize anything at all, and must report the existing snapshot.
|
||||
func TestRunStageWeights_IdempotentReStageIsNoop(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, weightsStore.PutWeights(ctx, selfHostModelID, "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
|
||||
msg, err := runStageWeights(ctx, explodingObjectStore{t: t}, weightsStore, t.TempDir(),
|
||||
selfHostModelID, "s3://models/llama-3.2-1b/", explodingMaterializer(t))
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, msg, "already staged")
|
||||
assert.Contains(t, msg, "snap-0001")
|
||||
|
||||
// The record must be untouched: still the same snapshot.
|
||||
entry, ok, err := weightsStore.GetWeights(ctx, selfHostModelID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "snap-0001", entry.SnapshotID)
|
||||
}
|
||||
|
||||
// TestRunStageWeights_ReStageDifferentSourceReplacesAndReportsPrevious
|
||||
// covers a re-stage from a new source: the KV entry is replaced with the new
|
||||
// source/snapshot, and the result reports the previous snapshot ID so an
|
||||
// operator can reclaim it.
|
||||
func TestRunStageWeights_ReStageDifferentSourceReplacesAndReportsPrevious(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, weightsStore.PutWeights(ctx, selfHostModelID, "s3://models/llama-3.2-1b-old/", "snap-old"))
|
||||
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
seedCompleteWeightsPrefix(t, store, "models", "llama-3.2-1b-new/")
|
||||
|
||||
materialize := func(context.Context, string, int64) (string, error) { return "snap-new", nil }
|
||||
|
||||
msg, err := runStageWeights(ctx, store, weightsStore, t.TempDir(),
|
||||
selfHostModelID, "s3://models/llama-3.2-1b-new/", materialize)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, msg, "snap-new")
|
||||
// The previous snapshot ID must appear in the result so an operator can reclaim it.
|
||||
assert.Contains(t, msg, "snap-old")
|
||||
|
||||
entry, ok, err := weightsStore.GetWeights(ctx, selfHostModelID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "s3://models/llama-3.2-1b-new/", entry.SourceURI)
|
||||
assert.Equal(t, "snap-new", entry.SnapshotID)
|
||||
}
|
||||
|
||||
// TestRunStageWeights_UnknownModelIDRefused covers refusal on a model ID
|
||||
// absent from the catalog entirely, before any store or materialize work.
|
||||
func TestRunStageWeights_UnknownModelIDRefused(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
|
||||
_, err := runStageWeights(context.Background(), explodingObjectStore{t: t}, weightsStore, t.TempDir(),
|
||||
"not-a-real-model", "s3://models/llama-3.2-1b/", explodingMaterializer(t))
|
||||
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "unknown model ID")
|
||||
|
||||
entries, listErr := weightsStore.ListWeights(context.Background())
|
||||
require.NoError(t, listErr)
|
||||
assert.Empty(t, entries, "refusal must not write a KV entry")
|
||||
}
|
||||
|
||||
// TestRunStageWeights_ProviderModelIDRefused covers refusal on a model ID
|
||||
// that exists in the catalog but is provider-served, not self-host. The
|
||||
// error must be distinguishable from the unknown-model-ID case, since
|
||||
// LookupServingSpec returns found and selfHost separately for exactly this
|
||||
// reason.
|
||||
func TestRunStageWeights_ProviderModelIDRefused(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
|
||||
_, err := runStageWeights(context.Background(), explodingObjectStore{t: t}, weightsStore, t.TempDir(),
|
||||
providerModelID, "s3://models/claude/", explodingMaterializer(t))
|
||||
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "provider-served")
|
||||
assert.NotContains(t, err.Error(), "unknown model ID", "must be distinguishable from the unknown-model-ID refusal")
|
||||
|
||||
entries, listErr := weightsStore.ListWeights(context.Background())
|
||||
require.NoError(t, listErr)
|
||||
assert.Empty(t, entries, "refusal must not write a KV entry")
|
||||
}
|
||||
|
||||
// TestRunStageWeights_MissingRequiredFilesRefusedBeforeMaterializing covers
|
||||
// refusal on an S3 prefix missing required Hugging Face files: validation
|
||||
// must fail before download or materialisation is ever attempted.
|
||||
func TestRunStageWeights_MissingRequiredFilesRefusedBeforeMaterializing(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
// Deliberately incomplete: no config.json, tokenizer files, or safetensors.
|
||||
putObject(t, store, "models", "incomplete/README.md", []byte("nothing useful"))
|
||||
|
||||
_, err := runStageWeights(context.Background(), store, weightsStore, t.TempDir(),
|
||||
selfHostModelID, "s3://models/incomplete/", explodingMaterializer(t))
|
||||
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "missing required file")
|
||||
|
||||
entries, listErr := weightsStore.ListWeights(context.Background())
|
||||
require.NoError(t, listErr)
|
||||
assert.Empty(t, entries, "refusal must not write a KV entry")
|
||||
}
|
||||
|
||||
// TestListWeightsOutput_NoModelsStaged covers the empty-catalog case.
|
||||
func TestListWeightsOutput_NoModelsStaged(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
|
||||
out, err := listWeightsOutput(context.Background(), weightsStore)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "No models staged.", out)
|
||||
}
|
||||
|
||||
// TestListWeightsOutput_ListsSeveralStagedModels covers the populated case:
|
||||
// every staged model's ID, source URI and snapshot ID must appear in the
|
||||
// rendered output.
|
||||
func TestListWeightsOutput_ListsSeveralStagedModels(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, weightsStore.PutWeights(ctx, selfHostModelID, "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
require.NoError(t, weightsStore.PutWeights(ctx, "amazon.titan-text-lite-v1", "s3://models/titan-text-lite/", "snap-0002"))
|
||||
|
||||
out, err := listWeightsOutput(ctx, weightsStore)
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, out, selfHostModelID)
|
||||
assert.Contains(t, out, "s3://models/llama-3.2-1b/")
|
||||
assert.Contains(t, out, "snap-0001")
|
||||
assert.Contains(t, out, "amazon.titan-text-lite-v1")
|
||||
assert.Contains(t, out, "s3://models/titan-text-lite/")
|
||||
assert.Contains(t, out, "snap-0002")
|
||||
}
|
||||
|
||||
// TestRemoveWeights_RoundTrip covers remove: it returns the dropped entry,
|
||||
// and afterward the model has no staged-weights record. It must not delete
|
||||
// anything outside the KV bucket -- removeWeights takes no object store at
|
||||
// all, so there is nothing for it to delete there.
|
||||
func TestRemoveWeights_RoundTrip(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
ctx := context.Background()
|
||||
require.NoError(t, weightsStore.PutWeights(ctx, selfHostModelID, "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
|
||||
entry, err := removeWeights(ctx, weightsStore, selfHostModelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "snap-0001", entry.SnapshotID)
|
||||
assert.Equal(t, "s3://models/llama-3.2-1b/", entry.SourceURI)
|
||||
|
||||
_, ok, err := weightsStore.GetWeights(ctx, selfHostModelID)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok, "KV entry must be dropped")
|
||||
}
|
||||
|
||||
// TestRemoveWeights_UnknownModelIsRefused covers remove against a model with
|
||||
// no staged entry: a clear error, not a silent no-op.
|
||||
func TestRemoveWeights_UnknownModelIsRefused(t *testing.T) {
|
||||
weightsStore := newWeightsStoreForTest(t)
|
||||
|
||||
_, err := removeWeights(context.Background(), weightsStore, selfHostModelID)
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "no staged weights entry")
|
||||
}
|
||||
|
||||
// TestBuildWeightsImage_SizesFileAndInvokesRunner covers the real,
|
||||
// environment-independent half of image building: the image file is
|
||||
// created and truncated to contentBytes plus 15% overhead, and the runner
|
||||
// is invoked with the right paths. mkfs.ext4 itself is faked so the test
|
||||
// does not require it on PATH.
|
||||
func TestBuildWeightsImage_SizesFileAndInvokesRunner(t *testing.T) {
|
||||
origRunner := mkfsExt4Runner
|
||||
t.Cleanup(func() { mkfsExt4Runner = origRunner })
|
||||
|
||||
var gotSrcDir, gotImagePath string
|
||||
mkfsExt4Runner = func(srcDir, imagePath string) error {
|
||||
gotSrcDir = srcDir
|
||||
gotImagePath = imagePath
|
||||
return nil
|
||||
}
|
||||
|
||||
srcDir := t.TempDir()
|
||||
imagePath := filepath.Join(t.TempDir(), "vol.img")
|
||||
contentBytes := int64(1_000_000_000) // 1 GB: 15% overhead dominates the 64 MiB floor.
|
||||
|
||||
require.NoError(t, buildWeightsImage(srcDir, imagePath, contentBytes))
|
||||
|
||||
assert.Equal(t, srcDir, gotSrcDir)
|
||||
assert.Equal(t, imagePath, gotImagePath)
|
||||
|
||||
stat, err := os.Stat(imagePath)
|
||||
require.NoError(t, err)
|
||||
wantSize := contentBytes + int64(float64(contentBytes)*0.15)
|
||||
assert.Equal(t, wantSize, stat.Size())
|
||||
}
|
||||
|
||||
// TestBuildWeightsImage_MinPaddingFloor covers small content: padding is
|
||||
// floored at 64 MiB rather than the (tiny) 15% overhead.
|
||||
func TestBuildWeightsImage_MinPaddingFloor(t *testing.T) {
|
||||
origRunner := mkfsExt4Runner
|
||||
t.Cleanup(func() { mkfsExt4Runner = origRunner })
|
||||
mkfsExt4Runner = func(string, string) error { return nil }
|
||||
|
||||
imagePath := filepath.Join(t.TempDir(), "vol.img")
|
||||
contentBytes := int64(1024) // 15% of this is far below the 64 MiB floor.
|
||||
|
||||
require.NoError(t, buildWeightsImage(t.TempDir(), imagePath, contentBytes))
|
||||
|
||||
stat, err := os.Stat(imagePath)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, contentBytes+64*1024*1024, stat.Size())
|
||||
}
|
||||
|
||||
// TestBuildWeightsImage_RunnerErrorPropagates covers mkfs.ext4 failing: the
|
||||
// error must surface, not be swallowed.
|
||||
func TestBuildWeightsImage_RunnerErrorPropagates(t *testing.T) {
|
||||
origRunner := mkfsExt4Runner
|
||||
t.Cleanup(func() { mkfsExt4Runner = origRunner })
|
||||
mkfsExt4Runner = func(string, string) error { return errors.New("mkfs boom") }
|
||||
|
||||
err := buildWeightsImage(t.TempDir(), filepath.Join(t.TempDir(), "vol.img"), 1024)
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "mkfs boom")
|
||||
}
|
||||
|
||||
// TestBuildWeightsImage_CreateFileErrorSkipsRunner covers a failure before
|
||||
// mkfs.ext4 would even run: an unwritable image path must fail without
|
||||
// invoking the runner.
|
||||
func TestBuildWeightsImage_CreateFileErrorSkipsRunner(t *testing.T) {
|
||||
origRunner := mkfsExt4Runner
|
||||
t.Cleanup(func() { mkfsExt4Runner = origRunner })
|
||||
called := false
|
||||
mkfsExt4Runner = func(string, string) error { called = true; return nil }
|
||||
|
||||
// A path inside a directory that doesn't exist cannot be created.
|
||||
badPath := filepath.Join(t.TempDir(), "missing-dir", "vol.img")
|
||||
err := buildWeightsImage(t.TempDir(), badPath, 1024)
|
||||
require.Error(t, err)
|
||||
assert.False(t, called, "runner must not be invoked when the image file cannot be created")
|
||||
}
|
||||
|
||||
// fakeClusterConfig returns a minimal single-node ClusterConfig sufficient
|
||||
// to drive the ochre weights Run wrappers: all nested config structs are
|
||||
// value types defaulting to the zero value, which every collaborator on the
|
||||
// refusal paths (LoadViperblockMasterKey, NewS3ObjectStoreFromConfig) treats
|
||||
// as "unset", not invalid.
|
||||
func fakeClusterConfig() *config.ClusterConfig {
|
||||
return &config.ClusterConfig{
|
||||
Node: "node1",
|
||||
Nodes: map[string]config.Config{"node1": {}},
|
||||
}
|
||||
}
|
||||
|
||||
// ochreExitPanic is the sentinel the ochreExit test stand-in panics with, so
|
||||
// a deferred recover can distinguish "the command called exit" from any
|
||||
// other panic escaping fn.
|
||||
type ochreExitPanic struct{ code int }
|
||||
|
||||
// withOchreExitCapture swaps ochreExit for a fake that records the exit code
|
||||
// and panics with ochreExitPanic instead of killing the test binary, runs
|
||||
// fn, and returns the recorded code (-1 if fn returned without exiting).
|
||||
// This lets a Run wrapper's real connect/validate/exit control flow be
|
||||
// exercised directly, the same way real operators experience it.
|
||||
func withOchreExitCapture(t *testing.T, fn func()) int {
|
||||
t.Helper()
|
||||
origExit := ochreExit
|
||||
t.Cleanup(func() { ochreExit = origExit })
|
||||
|
||||
code := -1
|
||||
ochreExit = func(c int) { code = c; panic(ochreExitPanic{c}) }
|
||||
|
||||
func() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
if _, ok := r.(ochreExitPanic); !ok {
|
||||
panic(r)
|
||||
}
|
||||
}
|
||||
}()
|
||||
fn()
|
||||
}()
|
||||
return code
|
||||
}
|
||||
|
||||
// captureStdout redirects os.Stdout for the duration of fn and returns what
|
||||
// was written, so wrapper tests can assert on the printed result without
|
||||
// relying on the underlying listWeightsOutput/removeWeights return value
|
||||
// alone.
|
||||
func captureStdout(t *testing.T, fn func()) string {
|
||||
t.Helper()
|
||||
r, w, err := os.Pipe()
|
||||
require.NoError(t, err)
|
||||
orig := os.Stdout
|
||||
os.Stdout = w //nolint:reassign // test-local stdout capture, restored below before this func returns
|
||||
|
||||
fn()
|
||||
|
||||
require.NoError(t, w.Close())
|
||||
os.Stdout = orig //nolint:reassign // restoring the real os.Stdout captured above
|
||||
var buf bytes.Buffer
|
||||
_, err = io.Copy(&buf, r)
|
||||
require.NoError(t, err)
|
||||
return buf.String()
|
||||
}
|
||||
|
||||
// stubConnect swaps loadConfigAndConnectFn for a fake returning cfg/nc/err,
|
||||
// restoring the real loadConfigAndConnect on cleanup.
|
||||
func stubConnect(t *testing.T, cfg *config.ClusterConfig, nc *nats.Conn, err error) {
|
||||
t.Helper()
|
||||
orig := loadConfigAndConnectFn
|
||||
t.Cleanup(func() { loadConfigAndConnectFn = orig })
|
||||
loadConfigAndConnectFn = func() (*config.ClusterConfig, *nats.Conn, error) { return cfg, nc, err }
|
||||
}
|
||||
|
||||
// newOchreWeightsStageTestCmd sets --model-id/--s3-uri/--tmp-dir on the real
|
||||
// ochreWeightsStageCmd, so tests drive the actual registered flags rather
|
||||
// than a hand-rolled stand-in.
|
||||
func newOchreWeightsStageTestCmd(t *testing.T, modelID, s3URI string) *cobra.Command {
|
||||
t.Helper()
|
||||
require.NoError(t, ochreWeightsStageCmd.Flags().Set("model-id", modelID))
|
||||
require.NoError(t, ochreWeightsStageCmd.Flags().Set("s3-uri", s3URI))
|
||||
require.NoError(t, ochreWeightsStageCmd.Flags().Set("tmp-dir", t.TempDir()))
|
||||
return ochreWeightsStageCmd
|
||||
}
|
||||
|
||||
func newOchreWeightsRemoveTestCmd(t *testing.T, modelID string) *cobra.Command {
|
||||
t.Helper()
|
||||
require.NoError(t, ochreWeightsRemoveCmd.Flags().Set("model-id", modelID))
|
||||
return ochreWeightsRemoveCmd
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsStage_ConnectFailureExits1 covers the wrapper's first
|
||||
// failure mode: it cannot even reach the catalog check without a cluster
|
||||
// connection.
|
||||
func TestRunOchreWeightsStage_ConnectFailureExits1(t *testing.T) {
|
||||
stubConnect(t, nil, nil, errors.New("dial failed"))
|
||||
|
||||
cmd := newOchreWeightsStageTestCmd(t, selfHostModelID, "s3://models/x/")
|
||||
code := withOchreExitCapture(t, func() { runOchreWeightsStage(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsStage_UnknownModelExits1 drives the real cobra Run
|
||||
// function end to end against an embedded JetStream server, confirming the
|
||||
// wrapper propagates runStageWeights's catalog refusal into a process exit.
|
||||
func TestRunOchreWeightsStage_UnknownModelExits1(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
stubConnect(t, fakeClusterConfig(), nc, nil)
|
||||
|
||||
cmd := newOchreWeightsStageTestCmd(t, "not-a-real-model", "s3://models/x/")
|
||||
code := withOchreExitCapture(t, func() { runOchreWeightsStage(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsStage_MasterKeyLoadFailureExits1 covers the wrapper's
|
||||
// own setup failing (an unreadable viperblock encryption key file) before
|
||||
// runStageWeights is ever called.
|
||||
func TestRunOchreWeightsStage_MasterKeyLoadFailureExits1(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
cfg := &config.ClusterConfig{
|
||||
Node: "node1",
|
||||
Nodes: map[string]config.Config{"node1": {
|
||||
Viperblock: config.ViperblockConfig{EncryptionKeyFile: filepath.Join(t.TempDir(), "missing-key")},
|
||||
}},
|
||||
}
|
||||
stubConnect(t, cfg, nc, nil)
|
||||
|
||||
cmd := newOchreWeightsStageTestCmd(t, selfHostModelID, "s3://models/x/")
|
||||
code := withOchreExitCapture(t, func() { runOchreWeightsStage(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsList_ConnectFailureExits1 mirrors the stage wrapper's
|
||||
// connect-failure coverage for list.
|
||||
func TestRunOchreWeightsList_ConnectFailureExits1(t *testing.T) {
|
||||
stubConnect(t, nil, nil, errors.New("dial failed"))
|
||||
|
||||
code := withOchreExitCapture(t, func() { runOchreWeightsList(nil, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsList_PrintsStagedModels drives the real Run function
|
||||
// end to end: it must print what a separately-constructed WeightsStore
|
||||
// against the same embedded JetStream server sees.
|
||||
func TestRunOchreWeightsList_PrintsStagedModels(t *testing.T) {
|
||||
_, nc, js := testutil.StartTestJetStream(t)
|
||||
stubConnect(t, fakeClusterConfig(), nc, nil)
|
||||
|
||||
seedStore := gateway_bedrock.NewWeightsStore(js, 1)
|
||||
require.NoError(t, seedStore.PutWeights(context.Background(), selfHostModelID, "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
|
||||
out := captureStdout(t, func() { runOchreWeightsList(nil, nil) })
|
||||
assert.Contains(t, out, selfHostModelID)
|
||||
assert.Contains(t, out, "snap-0001")
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsList_PrintsNoModelsStaged covers the empty case through
|
||||
// the full wrapper.
|
||||
func TestRunOchreWeightsList_PrintsNoModelsStaged(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
stubConnect(t, fakeClusterConfig(), nc, nil)
|
||||
|
||||
out := captureStdout(t, func() { runOchreWeightsList(nil, nil) })
|
||||
assert.Contains(t, out, "No models staged.")
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsRemove_ConnectFailureExits1 mirrors the stage wrapper's
|
||||
// connect-failure coverage for remove.
|
||||
func TestRunOchreWeightsRemove_ConnectFailureExits1(t *testing.T) {
|
||||
stubConnect(t, nil, nil, errors.New("dial failed"))
|
||||
|
||||
cmd := newOchreWeightsRemoveTestCmd(t, selfHostModelID)
|
||||
code := withOchreExitCapture(t, func() { runOchreWeightsRemove(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsRemove_UnknownModelExits1 drives the real cobra Run
|
||||
// function against an embedded JetStream server with nothing staged.
|
||||
func TestRunOchreWeightsRemove_UnknownModelExits1(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
stubConnect(t, fakeClusterConfig(), nc, nil)
|
||||
|
||||
cmd := newOchreWeightsRemoveTestCmd(t, selfHostModelID)
|
||||
code := withOchreExitCapture(t, func() { runOchreWeightsRemove(cmd, nil) })
|
||||
assert.Equal(t, 1, code)
|
||||
}
|
||||
|
||||
// TestRunOchreWeightsRemove_RoundTripPrintsSnapshotAndDropsEntry drives the
|
||||
// full wrapper: it must print the removed snapshot ID and leave the model
|
||||
// with no staged-weights entry afterward.
|
||||
func TestRunOchreWeightsRemove_RoundTripPrintsSnapshotAndDropsEntry(t *testing.T) {
|
||||
ns, nc, js := testutil.StartTestJetStream(t)
|
||||
stubConnect(t, fakeClusterConfig(), nc, nil)
|
||||
|
||||
seedStore := gateway_bedrock.NewWeightsStore(js, 1)
|
||||
require.NoError(t, seedStore.PutWeights(context.Background(), selfHostModelID, "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
|
||||
cmd := newOchreWeightsRemoveTestCmd(t, selfHostModelID)
|
||||
out := captureStdout(t, func() { runOchreWeightsRemove(cmd, nil) })
|
||||
assert.Contains(t, out, "snap-0001")
|
||||
|
||||
// runOchreWeightsRemove closed nc on return (mirroring real CLI usage: one
|
||||
// connection per invocation), so verify the drop over a fresh connection
|
||||
// to the same embedded server rather than reusing the closed one.
|
||||
verifyNC, err := nats.Connect(ns.ClientURL())
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { verifyNC.Close() })
|
||||
verifyStore := gateway_bedrock.NewWeightsStore(testutil.NewJetStream(t, verifyNC), 1)
|
||||
|
||||
_, ok, err := verifyStore.GetWeights(context.Background(), selfHostModelID)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok, "KV entry must be dropped")
|
||||
}
|
||||
|
||||
// failingObjectStore wraps a MemoryObjectStore and forces selected
|
||||
// operations to fail with a generic (non-NoSuchKey) error, covering error
|
||||
// branches a well-formed prefix never exercises: a transport/permission
|
||||
// failure, as opposed to an absent key.
|
||||
type failingObjectStore struct {
|
||||
*objectstore.MemoryObjectStore
|
||||
|
||||
failListObjects bool
|
||||
failGetObject bool
|
||||
// failHeadObjectKeys, if set, fails HeadObject only for these keys.
|
||||
failHeadObjectKeys map[string]bool
|
||||
}
|
||||
|
||||
func (f *failingObjectStore) ListObjectsV2(ctx context.Context, input *s3.ListObjectsV2Input) (*s3.ListObjectsV2Output, error) {
|
||||
if f.failListObjects {
|
||||
return nil, errors.New("list boom")
|
||||
}
|
||||
return f.MemoryObjectStore.ListObjectsV2(ctx, input)
|
||||
}
|
||||
|
||||
func (f *failingObjectStore) GetObject(ctx context.Context, input *s3.GetObjectInput) (*s3.GetObjectOutput, error) {
|
||||
if f.failGetObject {
|
||||
return nil, errors.New("get boom")
|
||||
}
|
||||
return f.MemoryObjectStore.GetObject(ctx, input)
|
||||
}
|
||||
|
||||
func (f *failingObjectStore) HeadObject(ctx context.Context, input *s3.HeadObjectInput) (*s3.HeadObjectOutput, error) {
|
||||
if f.failHeadObjectKeys[aws.StringValue(input.Key)] {
|
||||
return nil, errors.New("head boom")
|
||||
}
|
||||
return f.MemoryObjectStore.HeadObject(ctx, input)
|
||||
}
|
||||
|
||||
var _ objectstore.ObjectStore = (*failingObjectStore)(nil)
|
||||
|
||||
// pagingObjectStore returns ListObjectsV2 results across two pages linked by
|
||||
// a continuation token, covering listWeightsPrefix's pagination loop, which
|
||||
// MemoryObjectStore's single-page response never exercises.
|
||||
type pagingObjectStore struct {
|
||||
*objectstore.MemoryObjectStore
|
||||
|
||||
pages []*s3.Object
|
||||
calls int
|
||||
}
|
||||
|
||||
func (p *pagingObjectStore) ListObjectsV2(context.Context, *s3.ListObjectsV2Input) (*s3.ListObjectsV2Output, error) {
|
||||
obj := p.pages[p.calls]
|
||||
p.calls++
|
||||
out := &s3.ListObjectsV2Output{Contents: []*s3.Object{obj}}
|
||||
if p.calls < len(p.pages) {
|
||||
truncated := true
|
||||
token := "next"
|
||||
out.IsTruncated = &truncated
|
||||
out.NextContinuationToken = &token
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// TestListWeightsPrefix_PaginatesAcrossContinuationToken covers the
|
||||
// continuation-token loop: every page must be walked, not just the first.
|
||||
func TestListWeightsPrefix_PaginatesAcrossContinuationToken(t *testing.T) {
|
||||
store := &pagingObjectStore{
|
||||
MemoryObjectStore: objectstore.NewMemoryObjectStore(),
|
||||
pages: []*s3.Object{
|
||||
{Key: aws.String("a")},
|
||||
{Key: aws.String("b")},
|
||||
},
|
||||
}
|
||||
|
||||
var got []string
|
||||
err := listWeightsPrefix(context.Background(), store, "models", "prefix/", func(obj *s3.Object) error {
|
||||
got = append(got, aws.StringValue(obj.Key))
|
||||
return nil
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{"a", "b"}, got)
|
||||
assert.Equal(t, 2, store.calls)
|
||||
}
|
||||
|
||||
// TestListWeightsPrefix_ListObjectsErrorPropagates covers a transport error
|
||||
// from ListObjectsV2 itself, as opposed to a missing key.
|
||||
func TestListWeightsPrefix_ListObjectsErrorPropagates(t *testing.T) {
|
||||
store := &failingObjectStore{MemoryObjectStore: objectstore.NewMemoryObjectStore(), failListObjects: true}
|
||||
|
||||
err := listWeightsPrefix(context.Background(), store, "models", "prefix/", func(*s3.Object) error { return nil })
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "list boom")
|
||||
}
|
||||
|
||||
// TestListWeightsPrefix_CallbackErrorStopsWalk covers fn's error aborting
|
||||
// the walk immediately, rather than being swallowed.
|
||||
func TestListWeightsPrefix_CallbackErrorStopsWalk(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
putObject(t, store, "models", "prefix/a", nil)
|
||||
|
||||
err := listWeightsPrefix(context.Background(), store, "models", "prefix/", func(*s3.Object) error {
|
||||
return errors.New("callback boom")
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "callback boom")
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_RequiredFileHeadErrorPropagates covers a
|
||||
// transport error checking a fixed-name required file, distinct from that
|
||||
// file simply being absent (objectstore.IsNoSuchKeyError).
|
||||
func TestValidateWeightsPrefix_RequiredFileHeadErrorPropagates(t *testing.T) {
|
||||
store := &failingObjectStore{
|
||||
MemoryObjectStore: objectstore.NewMemoryObjectStore(),
|
||||
failHeadObjectKeys: map[string]bool{"llama-3.2-1b/config.json": true},
|
||||
}
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "head boom")
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_SafetensorsListErrorPropagates covers a
|
||||
// transport error listing for *.safetensors, distinct from none being
|
||||
// present.
|
||||
func TestValidateWeightsPrefix_SafetensorsListErrorPropagates(t *testing.T) {
|
||||
store := &failingObjectStore{MemoryObjectStore: objectstore.NewMemoryObjectStore(), failListObjects: true}
|
||||
for _, name := range requiredWeightsFiles {
|
||||
putObject(t, store.MemoryObjectStore, "models", "llama-3.2-1b/"+name, nil)
|
||||
}
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "list boom")
|
||||
}
|
||||
|
||||
// TestValidateWeightsPrefix_TokenizerHeadErrorPropagates covers a transport
|
||||
// error checking the tokenizer files, distinct from neither being present.
|
||||
func TestValidateWeightsPrefix_TokenizerHeadErrorPropagates(t *testing.T) {
|
||||
store := &failingObjectStore{
|
||||
MemoryObjectStore: objectstore.NewMemoryObjectStore(),
|
||||
failHeadObjectKeys: map[string]bool{
|
||||
"llama-3.2-1b/tokenizer.json": true,
|
||||
"llama-3.2-1b/tokenizer.model": true,
|
||||
},
|
||||
}
|
||||
for _, name := range requiredWeightsFiles {
|
||||
putObject(t, store.MemoryObjectStore, "models", "llama-3.2-1b/"+name, nil)
|
||||
}
|
||||
putObject(t, store.MemoryObjectStore, "models", "llama-3.2-1b/model.safetensors", nil)
|
||||
|
||||
err := validateWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/")
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "head boom")
|
||||
}
|
||||
|
||||
// TestDownloadWeightsPrefix_SkipsDirectoryMarkers covers predastore
|
||||
// directory-marker keys (ending in '/'): they must be skipped, not
|
||||
// downloaded as zero-byte files.
|
||||
func TestDownloadWeightsPrefix_SkipsDirectoryMarkers(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
putObject(t, store, "models", "llama-3.2-1b/subdir/", nil)
|
||||
putObject(t, store, "models", "llama-3.2-1b/config.json", []byte("{}"))
|
||||
|
||||
destDir := t.TempDir()
|
||||
total, err := downloadWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/", destDir)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(len("{}")), total)
|
||||
|
||||
entries, err := os.ReadDir(destDir)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, entries, 1)
|
||||
assert.Equal(t, "config.json", entries[0].Name())
|
||||
}
|
||||
|
||||
// TestDownloadWeightsPrefix_GetObjectErrorPropagates covers a download
|
||||
// failure surfacing through downloadWeightsPrefix's walk, not being
|
||||
// swallowed.
|
||||
func TestDownloadWeightsPrefix_GetObjectErrorPropagates(t *testing.T) {
|
||||
store := &failingObjectStore{MemoryObjectStore: objectstore.NewMemoryObjectStore(), failGetObject: true}
|
||||
putObject(t, store.MemoryObjectStore, "models", "llama-3.2-1b/config.json", []byte("{}"))
|
||||
|
||||
_, err := downloadWeightsPrefix(context.Background(), store, "models", "llama-3.2-1b/", t.TempDir())
|
||||
require.Error(t, err)
|
||||
assert.ErrorContains(t, err, "get boom")
|
||||
}
|
||||
|
||||
// TestDownloadObjectTo_OpenFileErrorPropagates covers the destination file
|
||||
// being uncreatable (parent directory missing): the error must surface
|
||||
// before any copy is attempted.
|
||||
func TestDownloadObjectTo_OpenFileErrorPropagates(t *testing.T) {
|
||||
store := objectstore.NewMemoryObjectStore()
|
||||
putObject(t, store, "models", "llama-3.2-1b/config.json", []byte("{}"))
|
||||
|
||||
badDest := filepath.Join(t.TempDir(), "missing-dir", "config.json")
|
||||
_, err := downloadObjectTo(context.Background(), store, "models", "llama-3.2-1b/config.json", badDest)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// lookupMkfsExt4 must find e2fsprogs in sbin, which a non-login shell's PATH
|
||||
// routinely omits. PATH is emptied so the fallback is what is under test.
|
||||
func TestLookupMkfsExt4_FallsBackToSbinDir(t *testing.T) {
|
||||
sbin := t.TempDir()
|
||||
require.NoError(t, os.WriteFile(filepath.Join(sbin, "mkfs.ext4"), []byte("#!/bin/sh\n"), 0o755))
|
||||
|
||||
t.Setenv("PATH", "")
|
||||
prev := mkfsExt4SbinDirs
|
||||
mkfsExt4SbinDirs = []string{sbin}
|
||||
t.Cleanup(func() { mkfsExt4SbinDirs = prev })
|
||||
|
||||
path, err := lookupMkfsExt4()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, filepath.Join(sbin, "mkfs.ext4"), path)
|
||||
}
|
||||
|
||||
func TestLookupMkfsExt4_AbsentEverywhereIsError(t *testing.T) {
|
||||
t.Setenv("PATH", "")
|
||||
prev := mkfsExt4SbinDirs
|
||||
mkfsExt4SbinDirs = []string{t.TempDir()}
|
||||
t.Cleanup(func() { mkfsExt4SbinDirs = prev })
|
||||
|
||||
_, err := lookupMkfsExt4()
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "install e2fsprogs")
|
||||
}
|
||||
|
||||
// A directory named mkfs.ext4 must not satisfy the lookup.
|
||||
func TestLookupMkfsExt4_IgnoresDirectory(t *testing.T) {
|
||||
sbin := t.TempDir()
|
||||
require.NoError(t, os.Mkdir(filepath.Join(sbin, "mkfs.ext4"), 0o755))
|
||||
|
||||
t.Setenv("PATH", "")
|
||||
prev := mkfsExt4SbinDirs
|
||||
mkfsExt4SbinDirs = []string{sbin}
|
||||
t.Cleanup(func() { mkfsExt4SbinDirs = prev })
|
||||
|
||||
_, err := lookupMkfsExt4()
|
||||
require.Error(t, err)
|
||||
}
|
||||
@@ -175,6 +175,14 @@ read.
|
||||
| `spx admin gpu mig enable` | `--profile <name>` (required, e.g. `1g.10gb`), `--gpu <pci-addr>` (optional, default: all MIG-capable GPUs) | Checks no GPU instances running (NATS) → discovers MIG-capable GPUs via `gpu.Discover()` (filtered by `--gpu` if set) → enables MIG mode on each target (`gpu.EnableMIGMode`) → lists available profiles (`gpu.ListProfiles`) and validates requested profile name → destroys any existing instances (`gpu.DestroyAllInstances`) → creates new instances filling GPU capacity (`gpu.CreateInstances`) → writes `mig_profile` to `spinifex.toml` via `admin.SetMIGProfile` → sends SIGHUP to `spinifex-daemon`. Must be run directly on the target host. |
|
||||
| `spx admin gpu mig disable` | `--gpu <pci-addr>` (optional, default: all MIG-capable GPUs) | Checks no GPU instances running (NATS) → discovers MIG-capable GPUs via `gpu.Discover()` (filtered by `--gpu` if set) → destroys all GPU instances (`gpu.DestroyAllInstances`) → disables MIG mode (`gpu.DisableMIGMode`) → clears `mig_profile` in `spinifex.toml` via `admin.SetMIGProfile` → sends SIGHUP to `spinifex-daemon`. Must be run directly on the target host. |
|
||||
|
||||
### Ochre Weights Management
|
||||
|
||||
| Command | Flags | Description |
|
||||
|---------|-------|-------------|
|
||||
| `spx admin ochre weights stage` | `--model-id` (required), `--s3-uri` (required, `s3://bucket/prefix/`), `--tmp-dir` (default: OS temp dir) | Stages a self-host model's weights for serving. Refuses first if `--model-id` is not a self-host catalog entry (unknown or provider-served), or if `bedrock-weights` already has this model staged from the same `--s3-uri` (no-op). Then validates the prefix holds `config.json`, `tokenizer_config.json`, `tokenizer.json`, `tokenizer.model` and at least one `*.safetensors` file, failing before downloading anything if any are missing. Downloads the prefix's objects from predastore into `--tmp-dir`, packs them into an ext4 image (`mkfs.ext4 -d`), imports it into a new viperblock volume via `v_utils.ImportDiskImage` with `AMIMetadata` left empty (plain volume, no AMI registration), snapshots the closed volume, and records the source URI and snapshot ID against `--model-id` in the `bedrock-weights` KV bucket. Re-staging a different `--s3-uri` replaces the entry and reports the previous snapshot ID for separate reclamation. |
|
||||
| `spx admin ochre weights list` | — | Lists every staged model with its source URI and snapshot ID, so an operator can see why a model is (or isn't) advertised via `ListFoundationModels`. |
|
||||
| `spx admin ochre weights remove` | `--model-id` (required) | Drops `--model-id`'s entry from `bedrock-weights`, hiding it from `ListFoundationModels` again. Refuses if the model has no staged entry. Never deletes the underlying snapshot or source S3 objects; reclaiming that storage is a separate, explicit act. |
|
||||
|
||||
### EKS Control-Plane Disaster Recovery
|
||||
|
||||
| Command | Flags | Description |
|
||||
|
||||
@@ -172,6 +172,37 @@ func lookupCatalogEntry(modelID string) (catalogEntry, bool) {
|
||||
return catalogEntry{}, false
|
||||
}
|
||||
|
||||
// ServingSpec is the subset of a self-host catalog entry needed both by
|
||||
// external callers (the 'ochre weights stage' CLI, validating a model ID
|
||||
// before staging weights) and by the daemon-side launcher (placing and
|
||||
// booting the serving VM), without exposing catalogEntry's AWS-shaped fields.
|
||||
type ServingSpec struct {
|
||||
ModelID string
|
||||
MinVRAMMiB int
|
||||
InstanceType string
|
||||
VLLMArgs []string
|
||||
}
|
||||
|
||||
// LookupServingSpec returns modelID's serving spec. found reports whether
|
||||
// modelID exists in the catalog at all; selfHost reports whether it is a
|
||||
// self-host entry (as opposed to a provider entry) when found is true. A
|
||||
// caller wanting the old found-and-servable behaviour checks found && selfHost.
|
||||
func LookupServingSpec(modelID string) (spec ServingSpec, found, selfHost bool) {
|
||||
entry, ok := lookupCatalogEntry(modelID)
|
||||
if !ok {
|
||||
return ServingSpec{}, false, false
|
||||
}
|
||||
if entry.Provider != tierSelfHost {
|
||||
return ServingSpec{}, true, false
|
||||
}
|
||||
return ServingSpec{
|
||||
ModelID: entry.ModelID,
|
||||
MinVRAMMiB: entry.MinVRAMMiB,
|
||||
InstanceType: entry.InstanceType,
|
||||
VLLMArgs: entry.VLLMArgs,
|
||||
}, true, true
|
||||
}
|
||||
|
||||
// ListFoundationModels returns the deployment-tiered catalog: self-host
|
||||
// entries where a weights snapshot resolves, provider entries where a
|
||||
// credential resolves.
|
||||
@@ -187,30 +218,6 @@ func ListFoundationModels(ctx context.Context, accountID string, resolver Creden
|
||||
return &bedrock.ListFoundationModelsOutput{ModelSummaries: summaries}, nil
|
||||
}
|
||||
|
||||
// ServingSpec is the subset of a self-host catalog entry the daemon-side
|
||||
// launcher needs to place and boot a serving VM. A separate type (rather than
|
||||
// exporting catalogEntry) keeps the launcher from depending on the AWS-shaped
|
||||
// summary/details fields this package alone renders.
|
||||
type ServingSpec struct {
|
||||
MinVRAMMiB int
|
||||
InstanceType string
|
||||
VLLMArgs []string
|
||||
}
|
||||
|
||||
// LookupServingSpec returns modelID's serving spec. ok is false for an
|
||||
// unknown model or a provider (non-self-host) entry, which has none.
|
||||
func LookupServingSpec(modelID string) (spec ServingSpec, ok bool) {
|
||||
entry, found := lookupCatalogEntry(modelID)
|
||||
if !found || entry.Provider != tierSelfHost {
|
||||
return ServingSpec{}, false
|
||||
}
|
||||
return ServingSpec{
|
||||
MinVRAMMiB: entry.MinVRAMMiB,
|
||||
InstanceType: entry.InstanceType,
|
||||
VLLMArgs: entry.VLLMArgs,
|
||||
}, true
|
||||
}
|
||||
|
||||
// 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,
|
||||
|
||||
@@ -132,3 +132,36 @@ func TestCatalog_SelfHostEntryCarriesServingSpec(t *testing.T) {
|
||||
assert.NotEmpty(t, entry.InstanceType)
|
||||
assert.NotEmpty(t, entry.VLLMArgs)
|
||||
}
|
||||
|
||||
// TestLookupServingSpec_SelfHostEntry covers 'ochre weights stage's happy
|
||||
// path: a known self-host model resolves found=true, selfHost=true, and a
|
||||
// non-empty serving spec.
|
||||
func TestLookupServingSpec_SelfHostEntry(t *testing.T) {
|
||||
spec, found, selfHost := LookupServingSpec("meta.llama3-2-1b-instruct-v1:0")
|
||||
require.True(t, found)
|
||||
require.True(t, selfHost)
|
||||
assert.Equal(t, "meta.llama3-2-1b-instruct-v1:0", spec.ModelID)
|
||||
assert.NotEmpty(t, spec.InstanceType)
|
||||
assert.Positive(t, spec.MinVRAMMiB)
|
||||
assert.NotEmpty(t, spec.VLLMArgs)
|
||||
}
|
||||
|
||||
// TestLookupServingSpec_ProviderEntry covers the refusal case 'stage' must
|
||||
// hit for a valid-but-wrong-tier model ID: found=true (it IS a catalog
|
||||
// entry), selfHost=false, so staging must be refused rather than proceed.
|
||||
func TestLookupServingSpec_ProviderEntry(t *testing.T) {
|
||||
spec, found, selfHost := LookupServingSpec("anthropic.claude-3-5-sonnet-20240620-v1:0")
|
||||
require.True(t, found)
|
||||
assert.False(t, selfHost)
|
||||
assert.Zero(t, spec)
|
||||
}
|
||||
|
||||
// TestLookupServingSpec_UnknownModel covers the other refusal case: a
|
||||
// typo'd or nonexistent model ID must report found=false, not merely
|
||||
// selfHost=false, so the CLI can give a precise error either way.
|
||||
func TestLookupServingSpec_UnknownModel(t *testing.T) {
|
||||
spec, found, selfHost := LookupServingSpec("not-a-real-model")
|
||||
assert.False(t, found)
|
||||
assert.False(t, selfHost)
|
||||
assert.Zero(t, spec)
|
||||
}
|
||||
|
||||
@@ -3,8 +3,10 @@ package gateway_bedrock
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"sync"
|
||||
|
||||
"github.com/mulgadc/spinifex/spinifex/kvutil"
|
||||
@@ -12,20 +14,29 @@ import (
|
||||
)
|
||||
|
||||
// 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.
|
||||
// model, this deployment's staged weights artifact. 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.
|
||||
// weightsKey returns the KV key for modelID's staged-weights record.
|
||||
// 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))
|
||||
}
|
||||
|
||||
// weightsRecord is the JSON value stored under weightsKey. SourceURI keeps
|
||||
// the record self-describing -- 'ochre weights stage' uses it to detect a
|
||||
// no-op re-stage, and 'ochre weights list' surfaces it so an operator can
|
||||
// see where a snapshot came from without a side channel.
|
||||
type weightsRecord struct {
|
||||
SnapshotID string `json:"snapshot_id"`
|
||||
SourceURI string `json:"source_uri"`
|
||||
}
|
||||
|
||||
// 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.
|
||||
@@ -33,8 +44,8 @@ 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.
|
||||
// WeightsStore resolves per-model weights records from the bedrock-weights
|
||||
// JetStream KV bucket.
|
||||
type WeightsStore struct {
|
||||
js jetstream.JetStream
|
||||
replicas int
|
||||
@@ -68,31 +79,122 @@ func (s *WeightsStore) bucket(ctx context.Context) (jetstream.KeyValue, error) {
|
||||
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) {
|
||||
// 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.
|
||||
func (s *WeightsStore) getRecord(ctx context.Context, modelID string) (weightsRecord, bool, error) {
|
||||
kv, err := s.bucket(ctx)
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
return weightsRecord{}, false, err
|
||||
}
|
||||
entry, err := kv.Get(ctx, weightsKey(modelID))
|
||||
switch {
|
||||
case err == nil:
|
||||
return string(entry.Value()), true, 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):
|
||||
return "", false, nil
|
||||
return weightsRecord{}, false, nil
|
||||
default:
|
||||
return "", false, fmt.Errorf("kv get weights snapshot for %s: %w", modelID, err)
|
||||
return weightsRecord{}, false, fmt.Errorf("kv get weights record 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 {
|
||||
// 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) {
|
||||
rec, ok, err := s.getRecord(ctx, modelID)
|
||||
if err != nil || !ok {
|
||||
return "", ok, err
|
||||
}
|
||||
return rec.SnapshotID, true, nil
|
||||
}
|
||||
|
||||
// GetWeights returns modelID's full staged-weights record (source URI and
|
||||
// snapshot ID), if one has been set. 'ochre weights stage' uses this to
|
||||
// decide idempotency and to report the snapshot a re-stage is replacing.
|
||||
func (s *WeightsStore) GetWeights(ctx context.Context, modelID string) (WeightsEntry, bool, error) {
|
||||
rec, ok, err := s.getRecord(ctx, modelID)
|
||||
if err != nil || !ok {
|
||||
return WeightsEntry{}, ok, err
|
||||
}
|
||||
return WeightsEntry{ModelID: modelID, SourceURI: rec.SourceURI, SnapshotID: rec.SnapshotID}, true, nil
|
||||
}
|
||||
|
||||
// PutWeights records modelID's staged weights artifact: the snapshot ID
|
||||
// endpoints COW-clone from, and the source S3 URI it was materialised from.
|
||||
func (s *WeightsStore) PutWeights(ctx context.Context, modelID, sourceURI, 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)
|
||||
value, err := json.Marshal(weightsRecord{SnapshotID: snapshotID, SourceURI: sourceURI})
|
||||
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)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// WeightsEntry is one staged model's KV record, decoded back from
|
||||
// weightsKey's base64url-encoded modelID for operator-facing listing.
|
||||
type WeightsEntry struct {
|
||||
ModelID string
|
||||
SourceURI string
|
||||
SnapshotID string
|
||||
}
|
||||
|
||||
// ListWeights returns every staged model and its record, sorted by model
|
||||
// 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.
|
||||
func (s *WeightsStore) ListWeights(ctx context.Context) ([]WeightsEntry, error) {
|
||||
kv, err := s.bucket(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
keys, err := kv.Keys(ctx)
|
||||
if err != nil {
|
||||
if errors.Is(err, jetstream.ErrNoKeysFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("kv list weights keys: %w", err)
|
||||
}
|
||||
|
||||
entries := make([]WeightsEntry, 0, len(keys))
|
||||
for _, key := range keys {
|
||||
modelID, err := base64.RawURLEncoding.DecodeString(key)
|
||||
if err != nil {
|
||||
// Not a key weightsKey wrote; skip rather than fail the whole list.
|
||||
continue
|
||||
}
|
||||
entry, err := kv.Get(ctx, key)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
var rec weightsRecord
|
||||
if err := json.Unmarshal(entry.Value(), &rec); err != nil {
|
||||
continue
|
||||
}
|
||||
entries = append(entries, WeightsEntry{ModelID: string(modelID), SourceURI: rec.SourceURI, SnapshotID: rec.SnapshotID})
|
||||
}
|
||||
sort.Slice(entries, func(i, j int) bool { return entries[i].ModelID < entries[j].ModelID })
|
||||
return entries, nil
|
||||
}
|
||||
|
||||
// DeleteWeights drops modelID's staged-weights KV entry only. The backing
|
||||
// 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)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -36,7 +36,7 @@ func TestWeightsStore_PutAndResolve_KV(t *testing.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", "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
|
||||
snapshotID, ok, err := store.Resolve(ctx, "meta.llama3-2-1b-instruct-v1:0")
|
||||
require.NoError(t, err)
|
||||
@@ -45,15 +45,15 @@ func TestWeightsStore_PutAndResolve_KV(t *testing.T) {
|
||||
}
|
||||
|
||||
// TestWeightsStore_PutOverwritesPrevious covers a re-stage: PutWeights
|
||||
// overwrites the prior snapshot ID in place rather than erroring or
|
||||
// accumulating history.
|
||||
// overwrites the prior record 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"))
|
||||
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
require.NoError(t, store.PutWeights(ctx, "meta.llama3-2-1b-instruct-v1:0", "s3://models/llama-3.2-1b-v2/", "snap-0002"))
|
||||
|
||||
snapshotID, ok, err := store.Resolve(ctx, "meta.llama3-2-1b-instruct-v1:0")
|
||||
require.NoError(t, err)
|
||||
@@ -61,6 +61,97 @@ func TestWeightsStore_PutOverwritesPrevious(t *testing.T) {
|
||||
assert.Equal(t, "snap-0002", snapshotID)
|
||||
}
|
||||
|
||||
// TestWeightsStore_GetWeights_Miss covers 'ochre weights stage' checking for
|
||||
// a prior record before materialising anything: a never-staged model
|
||||
// reports ok=false, not an error.
|
||||
func TestWeightsStore_GetWeights_Miss(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
|
||||
|
||||
entry, ok, err := store.GetWeights(context.Background(), "meta.llama3-2-1b-instruct-v1:0")
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok)
|
||||
assert.Zero(t, entry)
|
||||
}
|
||||
|
||||
// TestWeightsStore_GetWeights_ReturnsSourceURIAndSnapshot covers the
|
||||
// idempotency check stage relies on: the same source URI staged twice must
|
||||
// be detectable as a no-op, and a different source must surface the
|
||||
// snapshot it is about to replace.
|
||||
func TestWeightsStore_GetWeights_ReturnsSourceURIAndSnapshot(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", "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
|
||||
entry, ok, err := store.GetWeights(ctx, "meta.llama3-2-1b-instruct-v1:0")
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, WeightsEntry{
|
||||
ModelID: "meta.llama3-2-1b-instruct-v1:0",
|
||||
SourceURI: "s3://models/llama-3.2-1b/",
|
||||
SnapshotID: "snap-0001",
|
||||
}, entry)
|
||||
}
|
||||
|
||||
// TestWeightsStore_ListWeights_Empty covers the no-keys-yet case: an empty
|
||||
// bucket must report an empty list, not an error (jetstream.ErrNoKeysFound).
|
||||
func TestWeightsStore_ListWeights_Empty(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
|
||||
|
||||
entries, err := store.ListWeights(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, entries)
|
||||
}
|
||||
|
||||
// TestWeightsStore_ListWeights_ReturnsAllStaged covers the operator-facing
|
||||
// listing: every staged model comes back with its modelID decoded (not the
|
||||
// base64url KV key), its source URI and its snapshot ID, sorted by model ID.
|
||||
func TestWeightsStore_ListWeights_ReturnsAllStaged(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", "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
require.NoError(t, store.PutWeights(ctx, "amazon.titan-text-lite-v1", "s3://models/titan-text-lite/", "snap-0002"))
|
||||
|
||||
entries, err := store.ListWeights(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, entries, 2)
|
||||
assert.Equal(t, []WeightsEntry{
|
||||
{ModelID: "amazon.titan-text-lite-v1", SourceURI: "s3://models/titan-text-lite/", SnapshotID: "snap-0002"},
|
||||
{ModelID: "meta.llama3-2-1b-instruct-v1:0", SourceURI: "s3://models/llama-3.2-1b/", SnapshotID: "snap-0001"},
|
||||
}, entries)
|
||||
}
|
||||
|
||||
// TestWeightsStore_DeleteWeights_RemovesEntry covers the KV-only removal:
|
||||
// DeleteWeights drops the record and Resolve reports not-found afterwards.
|
||||
func TestWeightsStore_DeleteWeights_RemovesEntry(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", "s3://models/llama-3.2-1b/", "snap-0001"))
|
||||
require.NoError(t, store.DeleteWeights(ctx, "meta.llama3-2-1b-instruct-v1:0"))
|
||||
|
||||
_, ok, err := store.Resolve(ctx, "meta.llama3-2-1b-instruct-v1:0")
|
||||
require.NoError(t, err)
|
||||
assert.False(t, ok)
|
||||
}
|
||||
|
||||
// TestWeightsStore_DeleteWeights_NeverStagedIsNotAnError covers 'ochre
|
||||
// weights remove' against a model that was never staged: the CLI checks
|
||||
// GetWeights first to print a friendly message, but DeleteWeights itself
|
||||
// must not error on a missing key either.
|
||||
func TestWeightsStore_DeleteWeights_NeverStagedIsNotAnError(t *testing.T) {
|
||||
_, nc, _ := testutil.StartTestJetStream(t)
|
||||
store := NewWeightsStore(testutil.NewJetStream(t, nc), 1)
|
||||
|
||||
require.NoError(t, store.DeleteWeights(context.Background(), "meta.llama3-2-1b-instruct-v1:0"))
|
||||
}
|
||||
|
||||
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"))
|
||||
|
||||
@@ -125,8 +125,8 @@ func (s *Service) Ensure(ctx context.Context, in *EnsureEndpointInput, _ string)
|
||||
if in == nil || in.ModelID == "" {
|
||||
return nil, errors.New(awserrors.ErrorInvalidParameterValue)
|
||||
}
|
||||
spec, ok := gateway_bedrock.LookupServingSpec(in.ModelID)
|
||||
if !ok {
|
||||
spec, specFound, selfHost := gateway_bedrock.LookupServingSpec(in.ModelID)
|
||||
if !specFound || !selfHost {
|
||||
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user