Fix/create image version get version from call return

.changelog/44960.txt

@@ -0,0 +1,3 @@
+```release-note:bug-fix
+resource/aws_sagemaker_image_version: Fix race condition when creating multiple versions concurrently
+```

internal/service/sagemaker/image_version.go

@@ -6,6 +6,8 @@ package sagemaker
 import (
 	"context"
 	"log"
+	"strconv"
+	"strings"
 
 	"github.com/YakDriver/regexache"
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -163,11 +165,13 @@ func resourceImageVersionCreate(ctx context.Context, d *schema.ResourceData, met
 		input.Aliases = flex.ExpandStringValueSet(v.(*schema.Set))
 	}
 
-	if _, err := conn.CreateImageVersion(ctx, &input); err != nil {
+	result, err := conn.CreateImageVersion(ctx, &input)
+	if err != nil {
 		return sdkdiag.AppendErrorf(diags, "creating SageMaker AI Image Version %s: %s", name, err)
 	}
 
-	out, err := waitImageVersionCreated(ctx, conn, name)
+	version := extractVersionFromARN(aws.ToString(result.ImageVersionArn))
+	out, err := waitImageVersionCreatedByVersion(ctx, conn, name, version)
 	if err != nil {
 		return sdkdiag.AppendErrorf(diags, "waiting for SageMaker AI Image Version (%s) to be created: %s", d.Id(), err)
 	}
@@ -328,36 +332,6 @@ func resourceImageVersionDelete(ctx context.Context, d *schema.ResourceData, met
 	return diags
 }
 
-// findImageVersionByName is used immediately after creation to poll for status of
-// the most recently created image version
-//
-// findImageVersionByTwoPartKey should be used for all subsequent operations once the
-// version number is assigned.
-func findImageVersionByName(ctx context.Context, conn *sagemaker.Client, name string) (*sagemaker.DescribeImageVersionOutput, error) {
-	input := sagemaker.DescribeImageVersionInput{
-		ImageName: aws.String(name),
-	}
-
-	output, err := conn.DescribeImageVersion(ctx, &input)
-
-	if errs.IsAErrorMessageContains[*awstypes.ResourceNotFound](err, "does not exist") {
-		return nil, &retry.NotFoundError{
-			LastError:   err,
-			LastRequest: input,
-		}
-	}
-
-	if err != nil {
-		return nil, err
-	}
-
-	if output == nil {
-		return nil, tfresource.NewEmptyResultError(input)
-	}
-
-	return output, nil
-}
-
 // findImageVersionByTwoPartKey is used to fetch a specific version once a version number
 // is assigned
 func findImageVersionByTwoPartKey(ctx context.Context, conn *sagemaker.Client, name string, version int32) (*sagemaker.DescribeImageVersionOutput, error) {
@@ -412,6 +386,18 @@ func findImageVersionAliasesByTwoPartKey(ctx context.Context, conn *sagemaker.Cl
 	return output.SageMakerImageVersionAliases, nil
 }
 
+// extractVersionFromARN extracts the version number from an ImageVersionArn
+// ARN format: arn:aws:sagemaker:region:account:image-version/image-name/version
+func extractVersionFromARN(arn string) int32 {
+	parts := strings.Split(arn, "/")
+	if len(parts) >= 3 {
+		if version, err := strconv.ParseInt(parts[len(parts)-1], 10, 32); err == nil {
+			return int32(version)
+		}
+	}
+	return 0
+}
+
 // expandImageVersionResourceID wraps flex.ExpandResourceId and handles conversion
 // of the version portion to an integer
 func expandImageVersionResourceID(id string) (string, int32, error) {

internal/service/sagemaker/status.go

@@ -60,9 +60,9 @@ func statusImage(ctx context.Context, conn *sagemaker.Client, name string) retry
 	}
 }
 
-func statusImageVersionByName(ctx context.Context, conn *sagemaker.Client, name string) retry.StateRefreshFunc {
+func statusImageVersionByTwoPartKey(ctx context.Context, conn *sagemaker.Client, name string, version int32) retry.StateRefreshFunc {
 	return func() (any, string, error) {
-		output, err := findImageVersionByName(ctx, conn, name)
+		output, err := findImageVersionByTwoPartKey(ctx, conn, name, version)
 
 		if tfresource.NotFound(err) {
 			return nil, "", nil

internal/service/sagemaker/wait.go

@@ -213,11 +213,11 @@ func waitImageDeleted(ctx context.Context, conn *sagemaker.Client, name string)
 	return err
 }
 
-func waitImageVersionCreated(ctx context.Context, conn *sagemaker.Client, name string) (*sagemaker.DescribeImageVersionOutput, error) {
+func waitImageVersionCreatedByVersion(ctx context.Context, conn *sagemaker.Client, name string, version int32) (*sagemaker.DescribeImageVersionOutput, error) {
 	stateConf := &retry.StateChangeConf{
 		Pending: enum.Slice(awstypes.ImageVersionStatusCreating),
 		Target:  enum.Slice(awstypes.ImageVersionStatusCreated),
-		Refresh: statusImageVersionByName(ctx, conn, name),
+		Refresh: statusImageVersionByTwoPartKey(ctx, conn, name, version),
 		Timeout: imageVersionCreatedTimeout,
 	}