fix: register ochre weights snapshot with the EC2 control plane (#690)

weights stage created the snapshot at the viperblock layer only, writing
the block checkpoint and config.json but not the metadata.json that
CreateVolume resolves a SnapshotId through. The endpoint launcher clones
the weights volume via ec2.CreateVolume, so every launch failed with
InvalidSnapshot.NotFound against a snapshot that was present on disk.

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)
 	}

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")
+}