mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-19 19:01:42 +00:00
refactor: shared SelectNewestImage function
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user