diff --git a/cmd/spinifex/cmd/admin_ochre.go b/cmd/spinifex/cmd/admin_ochre.go index bb06e65dd..a8f8fb6de 100644 --- a/cmd/spinifex/cmd/admin_ochre.go +++ b/cmd/spinifex/cmd/admin_ochre.go @@ -18,6 +18,7 @@ import ( "github.com/mulgadc/spinifex/spinifex/admin" "github.com/mulgadc/spinifex/spinifex/config" gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" + handlers_ec2_snapshot "github.com/mulgadc/spinifex/spinifex/handlers/ec2/snapshot" "github.com/mulgadc/spinifex/spinifex/objectstore" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/mulgadc/viperblock/viperblock" @@ -311,7 +312,10 @@ func buildWeightsImage(srcDir, imagePath string, contentBytes int64) error { // 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) { +// +// az names the zone recorded on the snapshot for DescribeSnapshots; it is not +// used to resolve the clone. +func snapshotImportedWeightsVolume(s3Config *vbs3.S3Config, volumeID string, volumeSize uint64, walDir, az string, mkey *masterkey.Key) (string, error) { vbConfig := viperblock.VB{ VolumeName: volumeID, VolumeSize: volumeSize, @@ -350,9 +354,41 @@ func snapshotImportedWeightsVolume(s3Config *vbs3.S3Config, volumeID string, vol if _, err := vb.CreateSnapshot(snapshotID); err != nil { return "", fmt.Errorf("create snapshot: %w", err) } + + store := objectstore.NewS3ObjectStoreFromConfig(s3Config.Host, s3Config.Region, s3Config.AccessKey, s3Config.SecretKey) + if err := registerWeightsSnapshot(store, s3Config.Bucket, snapshotID, volumeID, volumeSize, az, mkey != nil); err != nil { + return "", err + } + return snapshotID, nil } +// registerWeightsSnapshot writes the EC2 control plane's half of the snapshot +// prefix. CreateSnapshot above writes only viperblock's half -- the block +// checkpoint and config.json -- but CreateVolume resolves a SnapshotId through +// metadata.json sitting alongside them, so without this the endpoint launcher +// cannot see the snapshot at all and fails with InvalidSnapshot.NotFound. +func registerWeightsSnapshot(store objectstore.ObjectStore, bucket, snapshotID, volumeID string, volumeSize uint64, az string, encrypted bool) error { + cfg := &handlers_ec2_snapshot.SnapshotConfig{ + SnapshotID: snapshotID, + VolumeID: volumeID, + // GiB, not bytes: CreateVolume compares this against a requested Size + // already in GiB, and would reject every clone if handed raw bytes. + VolumeSize: utils.SafeUint64ToInt64(volumeSize / bytesPerGiB), + State: "completed", + Progress: "100%", + StartTime: time.Now(), + Description: fmt.Sprintf("Ochre weights volume %s", volumeID), + Encrypted: encrypted, + OwnerID: utils.GlobalAccountID, + AvailabilityZone: az, + } + if err := handlers_ec2_snapshot.WriteSnapshotConfig(store, bucket, snapshotID, cfg); err != nil { + return fmt.Errorf("register snapshot metadata: %w", err) + } + return 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 @@ -514,7 +550,7 @@ func materializeWeightsVolume(node config.Config, downloadDir string, contentByt } fmt.Println("Snapshotting weights volume ...") - snapshotID, err := snapshotImportedWeightsVolume(&s3Config, volumeId, volumeBytes, tmpDir, mkey) + snapshotID, err := snapshotImportedWeightsVolume(&s3Config, volumeId, volumeBytes, tmpDir, node.AZ, mkey) if err != nil { return "", fmt.Errorf("could not snapshot weights volume: %w", err) } diff --git a/cmd/spinifex/cmd/admin_ochre_test.go b/cmd/spinifex/cmd/admin_ochre_test.go index 2185a2936..b6b9ce592 100644 --- a/cmd/spinifex/cmd/admin_ochre_test.go +++ b/cmd/spinifex/cmd/admin_ochre_test.go @@ -13,6 +13,7 @@ import ( "github.com/aws/aws-sdk-go/service/s3" "github.com/mulgadc/spinifex/spinifex/config" gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" + handlers_ec2_snapshot "github.com/mulgadc/spinifex/spinifex/handlers/ec2/snapshot" "github.com/mulgadc/spinifex/spinifex/objectstore" "github.com/mulgadc/spinifex/spinifex/testutil" "github.com/nats-io/nats.go" @@ -724,10 +725,18 @@ type failingObjectStore struct { failListObjects bool failGetObject bool + failPutObject bool // failHeadObjectKeys, if set, fails HeadObject only for these keys. failHeadObjectKeys map[string]bool } +func (f *failingObjectStore) PutObject(ctx context.Context, input *s3.PutObjectInput) (*s3.PutObjectOutput, error) { + if f.failPutObject { + return nil, errors.New("put boom") + } + return f.MemoryObjectStore.PutObject(ctx, input) +} + func (f *failingObjectStore) ListObjectsV2(ctx context.Context, input *s3.ListObjectsV2Input) (*s3.ListObjectsV2Output, error) { if f.failListObjects { return nil, errors.New("list boom") @@ -949,3 +958,45 @@ func TestLookupMkfsExt4_IgnoresDirectory(t *testing.T) { _, err := lookupMkfsExt4() require.Error(t, err) } + +// The launcher clones the weights snapshot through the EC2 control plane, so +// what stage writes must be readable by the same helper CreateVolume reads +// with, and carry the two fields it consumes: VolumeID (the viperblock source +// prefix) and VolumeSize (compared against a requested size, so GiB). +func TestRegisterWeightsSnapshot_WritesEC2ReadableMetadata(t *testing.T) { + store := objectstore.NewMemoryObjectStore() + const bucket = "predastore" + require.NoError(t, registerWeightsSnapshot(store, bucket, + "snap-vol-abc", "vol-abc", 12*bytesPerGiB, "ap-southeast-2a", true)) + + cfg, err := handlers_ec2_snapshot.ReadSnapshotConfig(store, bucket, "snap-vol-abc") + require.NoError(t, err) + assert.Equal(t, "snap-vol-abc", cfg.SnapshotID) + assert.Equal(t, "vol-abc", cfg.VolumeID) + assert.Equal(t, int64(12), cfg.VolumeSize) + assert.Equal(t, "completed", cfg.State) + assert.Equal(t, "ap-southeast-2a", cfg.AvailabilityZone) + assert.True(t, cfg.Encrypted) +} + +// An unencrypted volume must not be advertised as encrypted, which would let a +// consumer skip a key it actually needs. +func TestRegisterWeightsSnapshot_UnencryptedVolume(t *testing.T) { + store := objectstore.NewMemoryObjectStore() + require.NoError(t, registerWeightsSnapshot(store, "predastore", + "snap-vol-plain", "vol-plain", bytesPerGiB, "ap-southeast-2a", false)) + + cfg, err := handlers_ec2_snapshot.ReadSnapshotConfig(store, "predastore", "snap-vol-plain") + require.NoError(t, err) + assert.False(t, cfg.Encrypted) +} + +// A failed metadata write must surface, not leave stage reporting a snapshot +// ID the launcher will later reject as not found. +func TestRegisterWeightsSnapshot_WriteFailureSurfaces(t *testing.T) { + store := &failingObjectStore{MemoryObjectStore: objectstore.NewMemoryObjectStore(), failPutObject: true} + err := registerWeightsSnapshot(store, "predastore", + "snap-vol-abc", "vol-abc", bytesPerGiB, "ap-southeast-2a", false) + require.Error(t, err) + assert.Contains(t, err.Error(), "register snapshot metadata") +}