diff --git a/cmd/spinifex/cmd/admin.go b/cmd/spinifex/cmd/admin.go index 8f77189cb..862d85476 100644 --- a/cmd/spinifex/cmd/admin.go +++ b/cmd/spinifex/cmd/admin.go @@ -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 } diff --git a/cmd/spinifex/cmd/admin_ochre.go b/cmd/spinifex/cmd/admin_ochre.go new file mode 100644 index 000000000..bb06e65dd --- /dev/null +++ b/cmd/spinifex/cmd/admin_ochre.go @@ -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) +} diff --git a/cmd/spinifex/cmd/admin_ochre_test.go b/cmd/spinifex/cmd/admin_ochre_test.go new file mode 100644 index 000000000..2185a2936 --- /dev/null +++ b/cmd/spinifex/cmd/admin_ochre_test.go @@ -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) +} diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index 0fe9f08d4..2d664381e 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -175,6 +175,14 @@ read. | `spx admin gpu mig enable` | `--profile ` (required, e.g. `1g.10gb`), `--gpu ` (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 ` (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 | diff --git a/spinifex/gateway/bedrock/catalog.go b/spinifex/gateway/bedrock/catalog.go index 98d58d70c..af51ee4a4 100644 --- a/spinifex/gateway/bedrock/catalog.go +++ b/spinifex/gateway/bedrock/catalog.go @@ -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, diff --git a/spinifex/gateway/bedrock/catalog_test.go b/spinifex/gateway/bedrock/catalog_test.go index c5174624b..e39c21676 100644 --- a/spinifex/gateway/bedrock/catalog_test.go +++ b/spinifex/gateway/bedrock/catalog_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/weights.go b/spinifex/gateway/bedrock/weights.go index d508eba23..075418e8a 100644 --- a/spinifex/gateway/bedrock/weights.go +++ b/spinifex/gateway/bedrock/weights.go @@ -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 } diff --git a/spinifex/gateway/bedrock/weights_test.go b/spinifex/gateway/bedrock/weights_test.go index ebb391390..692485ba1 100644 --- a/spinifex/gateway/bedrock/weights_test.go +++ b/spinifex/gateway/bedrock/weights_test.go @@ -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")) diff --git a/spinifex/handlers/bedrock/service.go b/spinifex/handlers/bedrock/service.go index d699b8b69..f16e6c01d 100644 --- a/spinifex/handlers/bedrock/service.go +++ b/spinifex/handlers/bedrock/service.go @@ -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) }