feat(aws): add marketplace RHEL AI image support

cmd/mapt/cmd/aws/hosts/rhelai.go

@@ -64,6 +64,7 @@ func getRHELAICreate() *cobra.Command {
 					Version:          viper.GetString(params.RhelAIVersion),
 					Accelerator:      viper.GetString(params.RhelAIAccelerator),
 					CustomImage:      viper.GetString(params.RhelAICustomImage),
+					Marketplace:      viper.GetBool(params.RhelAIMarketplace),
 					ComputeRequest:   params.ComputeRequestArgs(),
 					Spot:             params.SpotArgs(),
 					Timeout:          viper.GetString(params.Timeout),
@@ -83,6 +84,7 @@ func getRHELAICreate() *cobra.Command {
 	flagSet.StringP(params.RhelAIVersion, "", params.RhelAIVersionDefault, params.RhelAIVersionDesc)
 	flagSet.StringP(params.RhelAIAccelerator, "", params.RhelAIAccelearatorDefault, params.RhelAIAccelearatorDesc)
 	flagSet.StringP(params.RhelAICustomImage, "", "", params.RhelAICustomImageDesc)
+	flagSet.Bool(params.RhelAIMarketplace, false, params.RhelAIMarketplaceDesc)
 	flagSet.StringP(params.RhelAIModel, "", "", params.RhelAIModelDesc)
 	flagSet.StringP(params.RhelAIHFToken, "", "", params.RhelAIHFTokenDesc)
 	flagSet.StringP(params.RhelAIAPIKey, "", "", params.RhelAIAPIKeyDesc)
@@ -133,7 +135,9 @@ func getRHELAIListVersions() *cobra.Command {
 			if err := viper.BindPFlags(cmd.Flags()); err != nil {
 				return err
 			}
-			versions, err := rhelai.ListVersions(cmd.Context(), viper.GetString(params.RhelAIAccelerator))
+			versions, err := rhelai.ListVersions(cmd.Context(),
+				viper.GetString(params.RhelAIAccelerator),
+				viper.GetBool(params.RhelAIMarketplace))
 			if err != nil {
 				return err
 			}
@@ -145,6 +149,7 @@ func getRHELAIListVersions() *cobra.Command {
 	}
 	flagSet := pflag.NewFlagSet(cmdRHELAIListVersions, pflag.ExitOnError)
 	flagSet.StringP(params.RhelAIAccelerator, "", params.RhelAIAccelearatorDefault, params.RhelAIAccelearatorDesc)
+	flagSet.Bool(params.RhelAIMarketplace, false, params.RhelAIMarketplaceDesc)
 	c.PersistentFlags().AddFlagSet(flagSet)
 	return c
 }

pkg/provider/aws/action/rhel-ai/constants.go

@@ -15,6 +15,9 @@ var (
 
 	// amiProduct     = "Red Hat Enterprise Linux"
 	amiProduct = "Linux/UNIX"
+
+	marketplaceOwner          = "679593333241"
+	marketplaceSupportedTypes = []string{"p4d.24xlarge", "p5.48xlarge", "g6e.48xlarge"}
 	amiV1Regex = "rhel-ai-nvidia-aws-%s*"
 	amiRegex   = "rhel-ai-%s-aws-%s*"
 	amiOwner   = "610952687893"
@@ -48,7 +51,7 @@ var (
 	outputUserPrivateKey = "ardPrivatekey"
 )
 
-func amiName(accelerator, version *string) string {
+func amiNameFromVersion(accelerator, version *string) string {
 	return util.If(strings.HasPrefix(*version, "1"),
 		fmt.Sprintf(amiV1Regex, *version),
 		fmt.Sprintf(amiRegex, *accelerator, *version))

pkg/provider/aws/action/rhel-ai/rhelai.go

@@ -26,6 +26,7 @@ import (
 	securityGroup "github.com/redhat-developer/mapt/pkg/provider/aws/services/ec2/security-group"
 	"github.com/redhat-developer/mapt/pkg/provider/util/command"
 	"github.com/redhat-developer/mapt/pkg/provider/util/output"
+	"github.com/redhat-developer/mapt/cmd/mapt/cmd/params"
 	apiRHELAI "github.com/redhat-developer/mapt/pkg/target/host/rhelai"
 	"github.com/redhat-developer/mapt/pkg/util"
 	"github.com/redhat-developer/mapt/pkg/util/logging"
@@ -38,6 +39,7 @@ type rhelAIRequest struct {
 	amiName          *string
 	arch             *string
 	spot             bool
+	marketplace      bool
 	timeout          *string
 	serviceEndpoints []string
 	allocationData   *allocation.AllocationResult
@@ -63,22 +65,45 @@ func (r *rhelAIRequest) validate() error {
 // If spot is enable it will run best spot option to get the best option to spin the machine
 // Then it will run the stack for windows dedicated host
 func Create(mCtxArgs *mc.ContextArgs, args *apiRHELAI.RHELAIArgs) (err error) {
+	if args.Marketplace && len(args.CustomImage) != 0 {
+		return fmt.Errorf("RHEL AI: --marketplace and --custom-image are mutually exclusive")
+	}
+	if args.Marketplace {
+		if err := validateMarketplaceComputeSizes(args.ComputeRequest.ComputeSizes); err != nil {
+			return err
+		}
+		if len(args.ComputeRequest.ComputeSizes) == 0 {
+			args.ComputeRequest.ComputeSizes = marketplaceSupportedTypes
+		}
+	}
 	// Create mapt Context
 	mCtx, err := mc.Init(mCtxArgs, aws.Provider())
 	if err != nil {
 		return err
 	}
 	// Compose request
-	amiName := amiName(&args.Accelerator, &args.Version)
-	if len(args.CustomImage) != 0 {
+	var amiName string
+	switch {
+	case args.Marketplace:
+		v := &args.Version
+		// When no explicit --version is given, ignore the default shared-AMI version
+		// and use a wildcard so the latest marketplace AMI is selected.
+		if args.Version == params.RhelAIVersionDefault {
+			v = nil
+		}
+		amiName = marketplaceAMIName(&args.Accelerator, v)
+	case len(args.CustomImage) != 0:
 		amiName = fmt.Sprintf("%s*", args.CustomImage)
+	default:
+		amiName = amiNameFromVersion(&args.Accelerator, &args.Version)
 	}
 	prefix := util.If(len(args.Prefix) > 0, args.Prefix, "main")
 	r := rhelAIRequest{
 		mCtx:             mCtx,
 		prefix:           &prefix,
 		amiName:          &amiName,
 		arch:             &args.Arch,
+		marketplace:      args.Marketplace,
 		timeout:          &args.Timeout,
 		serviceEndpoints: args.ServiceEndpoints,
 		diskSize:         args.ComputeRequest.DiskSize,
@@ -91,18 +116,30 @@ func Create(mCtxArgs *mc.ContextArgs, args *apiRHELAI.RHELAIArgs) (err error) {
 	if args.Spot != nil {
 		r.spot = args.Spot.Spot
 	}
-	r.allocationData, err = allocation.Allocation(mCtx,
-		&allocation.AllocationArgs{
-			Prefix:                &args.Prefix,
-			ComputeRequest:        args.ComputeRequest,
-			AMIProductDescription: &amiProduct,
-			Spot:                  args.Spot,
-		})
+	allocArgs := &allocation.AllocationArgs{
+		Prefix:                &args.Prefix,
+		ComputeRequest:        args.ComputeRequest,
+		AMIProductDescription: &amiProduct,
+		AMIName:               &amiName,
+		Spot:                  args.Spot,
+	}
+	if args.Marketplace {
+		owner := marketplaceOwner
+		allocArgs.AMIOwner = &owner
+		allocArgs.AMIPublic = true
+	}
+	r.allocationData, err = allocation.Allocation(mCtx, allocArgs)
 	if err != nil {
 		return err
 	}
-	if err = checkAMIExists(mCtx.Context(), &amiName, r.allocationData.Region, &amiArch); err != nil {
-		return err
+	if args.Marketplace {
+		if err = checkMarketplaceAMIExists(mCtx.Context(), &amiName, r.allocationData.Region, &amiArch); err != nil {
+			return err
+		}
+	} else {
+		if err = checkAMIExists(mCtx.Context(), &amiName, r.allocationData.Region, &amiArch); err != nil {
+			return err
+		}
 	}
 	return r.createMachine()
 }
@@ -136,9 +173,14 @@ func Destroy(mCtxArgs *mc.ContextArgs) error {
 const listVersionsRegion = "us-east-1"
 
 // ListVersions returns available RHEL AI version strings for the given accelerator,
-// sorted in ascending order. Versions are derived from AMI names in a reference
-// region (us-east-1) matching the pattern "rhel-ai-{accelerator}-aws-{version}*".
-func ListVersions(ctx context.Context, accelerator string) ([]string, error) {
+// sorted in ascending order. When marketplace is false, versions are derived from
+// AMI names in a reference region (us-east-1) matching "rhel-ai-{accelerator}-aws-{version}*".
+// When marketplace is true, versions are derived from marketplace AMI names
+// filtered by product code.
+func ListVersions(ctx context.Context, accelerator string, marketplace bool) ([]string, error) {
+	if marketplace {
+		return listMarketplaceVersions(ctx, accelerator)
+	}
 	acc := strings.ToLower(strings.TrimSpace(accelerator))
 	switch acc {
 	case "cuda", "rocm":
@@ -184,6 +226,35 @@ func ListVersions(ctx context.Context, accelerator string) ([]string, error) {
 	return versions, nil
 }
 
+func listMarketplaceVersions(ctx context.Context, accelerator string) ([]string, error) {
+	nameFilter := marketplaceAMIName(&accelerator, nil)
+	owner := marketplaceOwner
+	region := listVersionsRegion
+	images, err := data.ListAMIs(ctx, data.ImageRequest{
+		Name:   &nameFilter,
+		Arch:   &amiArch,
+		Owner:  &owner,
+		Region: &region,
+		Public: true,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("listing marketplace RHEL AI AMIs: %w", err)
+	}
+	seen := make(map[string]struct{})
+	for _, img := range images {
+		if img.Name == nil {
+			continue
+		}
+		seen[*img.Name] = struct{}{}
+	}
+	versions := make([]string, 0, len(seen))
+	for v := range seen {
+		versions = append(versions, v)
+	}
+	sort.Strings(versions)
+	return versions, nil
+}
+
 func (r *rhelAIRequest) createMachine() error {
 	cs := manager.Stack{
 		StackName:   r.mCtx.StackNameByProject(stackName),
@@ -211,12 +282,25 @@ func (r *rhelAIRequest) deploy(ctx *pulumi.Context) error {
 		return err
 	}
 	// Get AMI
-	ami, err := amiSVC.GetAMIByName(ctx,
-		*r.amiName,
-		[]string{amiOwner},
-		map[string]string{
-			"architecture": amiArch})
+	var ami *ec2.LookupAmiResult
+	var err error
+	if r.marketplace {
+		ami, err = amiSVC.GetAMIByName(ctx,
+			*r.amiName,
+			[]string{marketplaceOwner},
+			map[string]string{
+				"architecture": amiArch})
+	} else {
+		ami, err = amiSVC.GetAMIByName(ctx,
+			*r.amiName,
+			[]string{amiOwner},
+			map[string]string{
+				"architecture": amiArch})
+	}
 	if err != nil {
+		if r.marketplace && strings.Contains(err.Error(), "OptInRequired") {
+			return fmt.Errorf("RHEL AI marketplace: subscription required; accept terms at the AWS Marketplace console for RHEL AI before retrying\n%w", err)
+		}
 		return err
 	}
 	// Networking
@@ -415,3 +499,47 @@ func checkAMIExists(ctx context.Context, amiName, region, arch *string) error {
 	}
 	return nil
 }
+
+func checkMarketplaceAMIExists(ctx context.Context, amiName, region, arch *string) error {
+	owner := marketplaceOwner
+	isAMIOffered, _, err := data.IsAMIOffered(
+		ctx,
+		data.ImageRequest{
+			Name:   amiName,
+			Arch:   arch,
+			Region: region,
+			Owner:  &owner,
+			Public: true,
+		})
+	if err != nil {
+		return err
+	}
+	if !isAMIOffered {
+		return fmt.Errorf("marketplace RHEL AI AMI %q could not be found in region: %s", *amiName, *region)
+	}
+	return nil
+}
+
+func validateMarketplaceComputeSizes(sizes []string) error {
+	for _, s := range sizes {
+		supported := false
+		for _, ms := range marketplaceSupportedTypes {
+			if s == ms {
+				supported = true
+				break
+			}
+		}
+		if !supported {
+			return fmt.Errorf("instance type %q is not supported by AWS Marketplace RHEL AI; supported types: %s",
+				s, strings.Join(marketplaceSupportedTypes, ", "))
+		}
+	}
+	return nil
+}
+
+func marketplaceAMIName(accelerator, version *string) string {
+	if version == nil || len(*version) == 0 {
+		return fmt.Sprintf("rhel-ai-%s-*-prod-usfzkgvxvt65w", *accelerator)
+	}
+	return fmt.Sprintf("rhel-ai-%s-%s*-prod-usfzkgvxvt65w", *accelerator, *version)
+}

pkg/provider/aws/data/ami.go

@@ -24,6 +24,7 @@ type ImageRequest struct {
 	Name, Arch, Owner *string
 	Region            *string
 	BlockDeviceType   *string
+	Public            bool
 }
 
 const ERROR_NO_AMI = "no AMI"
@@ -62,12 +63,14 @@ func GetAMI(ctx context.Context, r ImageRequest) (*ImageInfo, error) {
 
 	if r.Owner != nil && len(*r.Owner) > 0 {
 		input.Owners = []string{*r.Owner}
-		aId, err := accountId(ctx)
-		if err != nil {
-			return nil, err
-		}
-		if *aId != *r.Owner {
-			input.ExecutableUsers = []string{"self"}
+		if !r.Public {
+			aId, err := accountId(ctx)
+			if err != nil {
+				return nil, err
+			}
+			if *aId != *r.Owner {
+				input.ExecutableUsers = []string{"self"}
+			}
 		}
 	}
 	result, err := client.DescribeImages(
@@ -134,12 +137,14 @@ func ListAMIs(ctx context.Context, r ImageRequest) ([]ec2Types.Image, error) {
 	}
 	if r.Owner != nil && len(*r.Owner) > 0 {
 		input.Owners = []string{*r.Owner}
-		aId, err := accountId(ctx)
-		if err != nil {
-			return nil, err
-		}
-		if *aId != *r.Owner {
-			input.ExecutableUsers = []string{"self"}
+		if !r.Public {
+			aId, err := accountId(ctx)
+			if err != nil {
+				return nil, err
+			}
+			if *aId != *r.Owner {
+				input.ExecutableUsers = []string{"self"}
+			}
 		}
 	}
 	result, err := client.DescribeImages(ctx, input)

pkg/provider/aws/data/spot.go

@@ -111,6 +111,8 @@ type SpotInfoArgs struct {
 	// AMI information
 	ProductDescription *string
 	AMIName, AMIArch   *string
+	AMIOwner           *string
+	AMIPublic          bool
 
 	ExcludedRegions       []string
 	SpotTolerance         *spot.Tolerance
@@ -146,7 +148,9 @@ func filterRegions(mCtx *mc.Context, args *SpotInfoArgs) ([]string, error) {
 					ImageRequest{
 						Name:   args.AMIName,
 						Arch:   args.AMIArch,
-						Region: &region})
+						Region: &region,
+						Owner:  args.AMIOwner,
+						Public: args.AMIPublic})
 				if err != nil {
 					if mCtx.Debug() {
 						logging.Warn(err.Error())

pkg/provider/aws/modules/allocation/allocation.go

@@ -16,8 +16,10 @@ type AllocationArgs struct {
 	ComputeRequest *cr.ComputeRequestArgs
 	Prefix,
 	AMIProductDescription,
-	AMIName *string
-	Spot *spotTypes.SpotArgs
+	AMIName,
+	AMIOwner *string
+	AMIPublic bool
+	Spot      *spotTypes.SpotArgs
 }
 
 type AllocationResult struct {
@@ -49,6 +51,10 @@ func Allocation(mCtx *mc.Context, args *AllocationArgs) (*AllocationResult, erro
 		if args.AMIProductDescription != nil {
 			sr.ProductDescription = *args.AMIProductDescription
 		}
+		if args.AMIOwner != nil {
+			sr.AMIOwner = *args.AMIOwner
+		}
+		sr.AMIPublic = args.AMIPublic
 		return allocationSpot(mCtx, sr)
 	}
 	return allocationOnDemand(mCtx, instancesTypes)