diff --git a/cmd/mapt/cmd/azure/hosts/rhelai.go b/cmd/mapt/cmd/azure/hosts/rhelai.go index 16c176a38..f3b424c1e 100644 --- a/cmd/mapt/cmd/azure/hosts/rhelai.go +++ b/cmd/mapt/cmd/azure/hosts/rhelai.go @@ -57,7 +57,7 @@ func getRHELAICreate() *cobra.Command { Prefix: "main", Version: viper.GetString(params.RhelAIVersion), Accelerator: viper.GetString(params.RhelAIAccelerator), - CustomAMI: viper.GetString(params.RhelAIAMICustom), + CustomImage: viper.GetString(params.RhelAIImageCustom), ComputeRequest: params.ComputeRequestArgs(), Spot: params.SpotArgs(), Timeout: viper.GetString(params.Timeout), @@ -69,7 +69,7 @@ func getRHELAICreate() *cobra.Command { flagSet.StringToStringP(params.Tags, "", nil, params.TagsDesc) flagSet.StringP(params.RhelAIVersion, "", params.RhelAIVersionDefault, params.RhelAIVersionDesc) flagSet.StringP(params.RhelAIAccelerator, "", params.RhelAIAccelearatorDefault, params.RhelAIAccelearatorDesc) - flagSet.StringP(params.RhelAIAMICustom, "", "", params.RhelAIAMICustomDesc) + flagSet.StringP(params.RhelAIImageCustom, "", "", params.RhelAIImageCustomDesc) flagSet.StringP(params.Timeout, "", "", params.TimeoutDesc) params.AddComputeRequestFlags(flagSet) params.AddSpotFlags(flagSet) diff --git a/cmd/mapt/cmd/params/params.go b/cmd/mapt/cmd/params/params.go index a5a8c598e..d69080a30 100644 --- a/cmd/mapt/cmd/params/params.go +++ b/cmd/mapt/cmd/params/params.go @@ -108,6 +108,8 @@ const ( RhelAIAccelearatorDefault string = "cuda" RhelAIAMICustom string = "custom-ami" RhelAIAMICustomDesc string = "custom AMI to spin RHEL AI OS" + RhelAIImageCustom string = "custom-image" + RhelAIImageCustomDesc string = "custom image ID to spin RHEL AI OS" // Serverless Timeout string = "timeout" diff --git a/pkg/provider/azure/action/rhel-ai/rhelai.go b/pkg/provider/azure/action/rhel-ai/rhelai.go index 13ae7060e..99becd420 100644 --- a/pkg/provider/azure/action/rhel-ai/rhelai.go +++ b/pkg/provider/azure/action/rhel-ai/rhelai.go @@ -22,17 +22,24 @@ const ( username = "azureuser" ) -func imageId(accelerator, version string) string { - iName := fmt.Sprintf(imageNameRegex, accelerator, version) - gName := strings.ReplaceAll(iName, "-", "_") +func imageIdFromName(imageName string) string { + gName := strings.ReplaceAll(imageName, "-", "_") return fmt.Sprintf(imageIdRegex, imageOwnerSubscriptionId, gName, - iName) + imageName) +} + +func imageId(accelerator, version string) string { + return imageIdFromName(fmt.Sprintf(imageNameRegex, accelerator, version)) } func Create(mCtxArgs *maptContext.ContextArgs, args *apiRHELAI.RHELAIArgs) (err error) { logging.Debug("Creating RHEL Server") + sharedImageID := imageId(args.Accelerator, args.Version) + if args.CustomImage != "" { + sharedImageID = imageIdFromName(args.CustomImage) + } azureLinuxRequest := &azureLinux.LinuxArgs{ Prefix: args.Prefix, @@ -40,7 +47,7 @@ func Create(mCtxArgs *maptContext.ContextArgs, args *apiRHELAI.RHELAIArgs) (err ComputeRequest: args.ComputeRequest, Spot: args.Spot, ImageRef: &data.ImageReference{ - SharedImageID: imageId(args.Accelerator, args.Version), + SharedImageID: sharedImageID, }, Username: username, ReadinessCommand: command.CommandPing} diff --git a/pkg/target/host/rhelai/api.go b/pkg/target/host/rhelai/api.go index 90e04bc96..8e6e9f493 100644 --- a/pkg/target/host/rhelai/api.go +++ b/pkg/target/host/rhelai/api.go @@ -10,6 +10,7 @@ type RHELAIArgs struct { Accelerator string Version string CustomAMI string + CustomImage string Arch string ComputeRequest *cr.ComputeRequestArgs Spot *spotTypes.SpotArgs