refactor: shared SelectNewestImage function

This commit is contained in:
Julian Sommer
2026-08-25 10:43:16 +10:00
parent 0d6676ecc8
commit e26fcdc7ac
6 changed files with 152 additions and 107 deletions
+2 -13
View File
@@ -43,19 +43,8 @@ func resolveServingAMI(ctx context.Context, imgSvc amiResolver) (string, error)
return "", ErrServingAMINotFound
}
var newestID, newestCreated string
matches := 0
for _, image := range out.Images {
if image == nil || aws.StringValue(image.ImageId) == "" {
continue
}
matches++
created := aws.StringValue(image.CreationDate)
if newestID == "" || created > newestCreated {
newestID = aws.StringValue(image.ImageId)
newestCreated = created
}
}
// Serving images are GPU-tagged by design, so nothing is excluded here.
newestID, _, matches := utils.SelectNewestImage(out.Images, "")
if newestID == "" {
return "", ErrServingAMINotFound
}
+7 -31
View File
@@ -11,6 +11,7 @@ import (
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/instancetypes"
"github.com/mulgadc/spinifex/spinifex/tags"
"github.com/mulgadc/spinifex/spinifex/utils"
)
const (
@@ -187,16 +188,6 @@ func lookupECSGPUNodeAMI(ctx context.Context, amiSvc ecsImageResolver, accountID
"ecs: multiple GPU AMIs match managed-by=ecs+gpu-vendor; using newest")
}
// hasTagKey reports whether img carries a tag with the given key.
func hasTagKey(img *ec2.Image, key string) bool {
for _, t := range img.Tags {
if aws.StringValue(t.Key) == key {
return true
}
}
return false
}
// resolveNewestAMI runs DescribeImages with filters and returns the newest
// (by CreationDate) matching image ID. describeErrCtx prefixes a DescribeImages
// failure; notFoundDesc/warnMsg describe the filter for the not-found error and
@@ -207,28 +198,13 @@ func resolveNewestAMI(ctx context.Context, amiSvc ecsImageResolver, accountID st
return "", fmt.Errorf("%s: %w", describeErrCtx, err)
}
var (
newestID string
newestCreated string
matches int
)
for _, img := range out.Images {
if img == nil || img.ImageId == nil || *img.ImageId == "" {
continue
}
// The GPU node AMI also carries managed-by=ecs; DescribeImages filters
// have no negation, so exclude gpu-vendor-tagged images client-side from
// the non-GPU lookup or a newer GPU AMI would hijack ordinary instances.
if excludeGPUTagged && hasTagKey(img, tags.GPUVendorKey) {
continue
}
matches++
// CreationDate is a fixed-width RFC3339 timestamp, so lexicographic
// comparison orders it correctly without parsing.
if created := aws.StringValue(img.CreationDate); newestID == "" || created > newestCreated {
newestID, newestCreated = *img.ImageId, created
}
// The GPU node AMI also carries managed-by=ecs, so the non-GPU lookup has to
// exclude gpu-vendor-tagged images or a newer GPU AMI would hijack ordinary instances.
excludeKey := ""
if excludeGPUTagged {
excludeKey = tags.GPUVendorKey
}
newestID, newestCreated, matches := utils.SelectNewestImage(out.Images, excludeKey)
if newestID == "" {
return "", fmt.Errorf("%w (%s, account %s)", notFound, notFoundDesc, accountID)
}
+7 -31
View File
@@ -16,6 +16,7 @@ import (
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/handlers/sysinstance"
"github.com/mulgadc/spinifex/spinifex/tags"
"github.com/mulgadc/spinifex/spinifex/utils"
)
// ErrEKSServerAMINotFound is returned when no AMI with the EKS managed-by tag
@@ -376,16 +377,6 @@ func lookupEKSGPUNodeAMI(ctx context.Context, amiSvc k3sAMIResolver, accountID,
"eks: multiple GPU AMIs match managed-by=eks+gpu-vendor; using newest")
}
// hasTagKey reports whether img carries a tag with the given key.
func hasTagKey(img *ec2.Image, key string) bool {
for _, t := range img.Tags {
if aws.StringValue(t.Key) == key {
return true
}
}
return false
}
// resolveNewestAMI runs DescribeImages with filters and returns the newest
// (by CreationDate) matching image ID. describeErrCtx prefixes a DescribeImages
// failure; notFoundDesc/warnMsg describe the filter for the not-found error and
@@ -396,28 +387,13 @@ func resolveNewestAMI(ctx context.Context, amiSvc k3sAMIResolver, accountID stri
return "", fmt.Errorf("%s: %w", describeErrCtx, err)
}
var (
newestID string
newestCreated string
matches int
)
for _, img := range out.Images {
if img == nil || img.ImageId == nil || *img.ImageId == "" {
continue
}
// The GPU node AMI also carries managed-by=eks; DescribeImages filters
// have no negation, so exclude gpu-vendor-tagged images client-side from
// the non-GPU lookup or a newer GPU AMI would hijack ordinary nodes.
if excludeGPUTagged && hasTagKey(img, tags.GPUVendorKey) {
continue
}
matches++
// CreationDate is a fixed-width RFC3339 timestamp, so lexicographic
// comparison orders it correctly without parsing.
if created := aws.StringValue(img.CreationDate); newestID == "" || created > newestCreated {
newestID, newestCreated = *img.ImageId, created
}
// The GPU node AMI also carries managed-by=eks, so the non-GPU lookup has to
// exclude gpu-vendor-tagged images or a newer GPU AMI would hijack ordinary nodes.
excludeKey := ""
if excludeGPUTagged {
excludeKey = tags.GPUVendorKey
}
newestID, newestCreated, matches := utils.SelectNewestImage(out.Images, excludeKey)
if newestID == "" {
return "", fmt.Errorf("%w (%s, account %s)", notFound, notFoundDesc, accountID)
}
+4 -31
View File
@@ -539,27 +539,10 @@ func resolveEngineAMI(ctx context.Context, amiSvc launchAMIResolver, engine, ver
return "", engineAMINotFound(engine, version)
}
// Several builds of one engine version can be registered; skip malformed
// catalog entries and select the most recently imported usable image.
var newestID, newestCreated string
matches := 0
for _, image := range out.Images {
if image == nil || aws.StringValue(image.ImageId) == "" {
continue
}
// A GPU engine build carries the same engine tags, and DescribeImages
// filters have no negation, so drop it here or a newer GPU image would
// hijack an ordinary instance.
if hasTagKey(image, tags.GPUVendorKey) {
continue
}
matches++
created := aws.StringValue(image.CreationDate)
if newestID == "" || created > newestCreated {
newestID = aws.StringValue(image.ImageId)
newestCreated = created
}
}
// Several builds of one engine version can be registered; select the most
// recently imported usable image. A GPU engine build carries the same engine
// tags, so it is excluded or a newer GPU image would hijack an ordinary instance.
newestID, _, matches := utils.SelectNewestImage(out.Images, tags.GPUVendorKey)
if newestID == "" {
return "", engineAMINotFound(engine, version)
}
@@ -570,16 +553,6 @@ func resolveEngineAMI(ctx context.Context, amiSvc launchAMIResolver, engine, ver
return newestID, nil
}
// hasTagKey reports whether img carries a tag with the given key.
func hasTagKey(img *ec2.Image, key string) bool {
for _, t := range img.Tags {
if aws.StringValue(t.Key) == key {
return true
}
}
return false
}
// Only the node running the VM subscribes the per-instance ec2.cmd subject, so
// the command routes itself and the caller need not know where the VM landed.
type natsVolumeAttacher struct {
+37
View File
@@ -24,6 +24,9 @@ import (
"runtime"
"strings"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
)
var ErrQCOWDetected = errors.New("qcow format detected")
@@ -237,6 +240,40 @@ func hashImageFile(imagePath string, hasher hash.Hash) (string, error) {
return hex.EncodeToString(hasher.Sum(nil)), nil
}
// SelectNewestImage returns the newest image ID among images, its CreationDate,
// and how many images were considered. Nil images, images with an empty ImageId
// and nil tag entries are skipped; a non-empty excludeTagKey drops images
// carrying that tag key. Returns an empty ID when nothing matches.
func SelectNewestImage(images []*ec2.Image, excludeTagKey string) (imageID, created string, matches int) {
for _, img := range images {
if img == nil || aws.StringValue(img.ImageId) == "" {
continue
}
// DescribeImages filters have no negation, so callers that must reject a
// tag (e.g. gpu-vendor) can only do it client-side, here.
if excludeTagKey != "" && imageHasTagKey(img, excludeTagKey) {
continue
}
matches++
// CreationDate is a fixed-width RFC3339 timestamp, so lexicographic
// comparison orders it correctly without parsing.
if c := aws.StringValue(img.CreationDate); imageID == "" || c > created {
imageID, created = aws.StringValue(img.ImageId), c
}
}
return imageID, created, matches
}
// imageHasTagKey reports whether img carries a tag with the given key.
func imageHasTagKey(img *ec2.Image, key string) bool {
for _, t := range img.Tags {
if t != nil && aws.StringValue(t.Key) == key {
return true
}
}
return false
}
type Images struct {
Name string `json:"name"`
Description string `json:"description"`
+95 -1
View File
@@ -1,6 +1,11 @@
package utils
import "testing"
import (
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
)
// resolveServingAMI filters purely on these two tags and never on the image
// name, so a typo in either hides the serving image from Ochre and every
@@ -37,3 +42,92 @@ func TestAvailableImages_VLLMServingCarriesBedrockTags(t *testing.T) {
t.Error("URL and Checksum must both be set or the import has nothing to fetch and verify")
}
}
func img(id, created string, tagKeys ...string) *ec2.Image {
image := &ec2.Image{ImageId: aws.String(id), CreationDate: aws.String(created)}
for _, k := range tagKeys {
if k == "" {
image.Tags = append(image.Tags, nil)
continue
}
image.Tags = append(image.Tags, &ec2.Tag{Key: aws.String(k), Value: aws.String("v")})
}
return image
}
func TestSelectNewestImage(t *testing.T) {
tests := []struct {
name string
images []*ec2.Image
excludeKey string
wantID string
wantCreated string
wantMatches int
}{
{
name: "empty input",
},
{
name: "all entries nil",
images: []*ec2.Image{nil, nil},
},
{
name: "empty image id skipped",
images: []*ec2.Image{img("", "2026-01-01T00:00:00.000Z"), img("ami-a", "2025-01-01T00:00:00.000Z")},
// The empty-ID entry is not a match, so no multi-match warning fires.
wantID: "ami-a", wantCreated: "2025-01-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "newest creation date wins regardless of order",
images: []*ec2.Image{img("ami-old", "2026-01-01T00:00:00.000Z"), img("ami-new", "2026-08-01T00:00:00.000Z"), img("ami-mid", "2026-04-01T00:00:00.000Z")},
wantID: "ami-new", wantCreated: "2026-08-01T00:00:00.000Z", wantMatches: 3,
},
{
name: "identical dates keep the first match",
images: []*ec2.Image{img("ami-a", "2026-01-01T00:00:00.000Z"), img("ami-b", "2026-01-01T00:00:00.000Z")},
wantID: "ami-a", wantCreated: "2026-01-01T00:00:00.000Z", wantMatches: 2,
},
{
name: "exclude tag key drops tagged images",
images: []*ec2.Image{img("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor"), img("ami-cpu", "2026-01-01T00:00:00.000Z")},
excludeKey: "gpu-vendor",
wantID: "ami-cpu", wantCreated: "2026-01-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "empty exclude key keeps tagged images",
images: []*ec2.Image{img("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor"), img("ami-cpu", "2026-01-01T00:00:00.000Z")},
wantID: "ami-gpu", wantCreated: "2026-08-01T00:00:00.000Z", wantMatches: 2,
},
{
name: "all images excluded",
images: []*ec2.Image{img("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor")},
excludeKey: "gpu-vendor",
},
{
name: "nil tag entry does not panic",
images: []*ec2.Image{img("ami-a", "2026-01-01T00:00:00.000Z", "", "gpu-vendor"), img("ami-b", "2026-02-01T00:00:00.000Z", "")},
excludeKey: "gpu-vendor",
wantID: "ami-b", wantCreated: "2026-02-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "missing creation date sorts oldest",
images: []*ec2.Image{&ec2.Image{ImageId: aws.String("ami-nodate")}, img("ami-dated", "2020-01-01T00:00:00.000Z")},
wantID: "ami-dated", wantCreated: "2020-01-01T00:00:00.000Z", wantMatches: 2,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gotID, gotCreated, gotMatches := SelectNewestImage(tt.images, tt.excludeKey)
if gotID != tt.wantID {
t.Errorf("imageID = %q, want %q", gotID, tt.wantID)
}
if gotCreated != tt.wantCreated {
t.Errorf("created = %q, want %q", gotCreated, tt.wantCreated)
}
if gotMatches != tt.wantMatches {
t.Errorf("matches = %d, want %d", gotMatches, tt.wantMatches)
}
})
}
}