From e26fcdc7aca9a783c4abae5dd920bc84fab583af Mon Sep 17 00:00:00 2001 From: Julian Sommer <115516918+juliansommer@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:43:16 +1000 Subject: [PATCH] refactor: shared SelectNewestImage function --- spinifex/handlers/bedrock/ami.go | 15 +--- spinifex/handlers/ecs/capacity.go | 38 ++-------- spinifex/handlers/eks/k3s_server_vm.go | 38 ++-------- spinifex/handlers/rds/launch.go | 35 ++-------- spinifex/utils/images.go | 37 ++++++++++ spinifex/utils/images_test.go | 96 +++++++++++++++++++++++++- 6 files changed, 152 insertions(+), 107 deletions(-) diff --git a/spinifex/handlers/bedrock/ami.go b/spinifex/handlers/bedrock/ami.go index a186f3f03..32885186f 100644 --- a/spinifex/handlers/bedrock/ami.go +++ b/spinifex/handlers/bedrock/ami.go @@ -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 } diff --git a/spinifex/handlers/ecs/capacity.go b/spinifex/handlers/ecs/capacity.go index 651a3f06f..1c7e9afe8 100644 --- a/spinifex/handlers/ecs/capacity.go +++ b/spinifex/handlers/ecs/capacity.go @@ -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) } diff --git a/spinifex/handlers/eks/k3s_server_vm.go b/spinifex/handlers/eks/k3s_server_vm.go index 9b94f3bc4..9d286ac09 100644 --- a/spinifex/handlers/eks/k3s_server_vm.go +++ b/spinifex/handlers/eks/k3s_server_vm.go @@ -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) } diff --git a/spinifex/handlers/rds/launch.go b/spinifex/handlers/rds/launch.go index 724496de2..cae61ae59 100644 --- a/spinifex/handlers/rds/launch.go +++ b/spinifex/handlers/rds/launch.go @@ -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 { diff --git a/spinifex/utils/images.go b/spinifex/utils/images.go index c10f4d750..31c6e50e5 100644 --- a/spinifex/utils/images.go +++ b/spinifex/utils/images.go @@ -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"` diff --git a/spinifex/utils/images_test.go b/spinifex/utils/images_test.go index 21a8dfa8b..441458987 100644 --- a/spinifex/utils/images_test.go +++ b/spinifex/utils/images_test.go @@ -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) + } + }) + } +}