rhelai.go

package rhelai

import (
	"context"
	"fmt"
	"sort"
	"strings"

	"github.com/go-playground/validator/v10"
	"github.com/pulumi/pulumi-aws/sdk/v7/go/aws/ec2"
	"github.com/pulumi/pulumi/sdk/v3/go/auto"
	"github.com/pulumi/pulumi/sdk/v3/go/pulumi"
	"github.com/redhat-developer/mapt/pkg/manager"
	mc "github.com/redhat-developer/mapt/pkg/manager/context"
	infra "github.com/redhat-developer/mapt/pkg/provider"
	"github.com/redhat-developer/mapt/pkg/provider/aws"
	awsConstants "github.com/redhat-developer/mapt/pkg/provider/aws/constants"
	"github.com/redhat-developer/mapt/pkg/provider/aws/data"
	"github.com/redhat-developer/mapt/pkg/provider/aws/modules/allocation"
	"github.com/redhat-developer/mapt/pkg/provider/aws/modules/ec2/compute"
	"github.com/redhat-developer/mapt/pkg/provider/aws/modules/network"
	"github.com/redhat-developer/mapt/pkg/provider/aws/modules/serverless"
	"github.com/redhat-developer/mapt/pkg/provider/aws/modules/spot"
	amiSVC "github.com/redhat-developer/mapt/pkg/provider/aws/services/ec2/ami"
	"github.com/redhat-developer/mapt/pkg/provider/aws/services/ec2/keypair"
	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"
	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"
	resourcesUtil "github.com/redhat-developer/mapt/pkg/util/resources"
)

type rhelAIRequest struct {
	mCtx             *mc.Context
	prefix           *string
	amiName          *string
	arch             *string
	spot             bool
	timeout          *string
	serviceEndpoints []string
	vpcID            *string
	allocationData   *allocation.AllocationResult
	diskSize         *int
	model            *string
	hfToken          *string
	apiKey           *string
	autoStart        bool
	vllmExtraArgs    *string
	exposePorts      []int
}

func (r *rhelAIRequest) validate() error {
	v := validator.New(validator.WithRequiredStructEnabled())
	err := v.Var(r.mCtx, "required")
	if err != nil {
		return err
	}
	return v.Struct(r)
}

// Create orchestrate 2 stacks:
// 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) {
	// 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 {
		amiName = fmt.Sprintf("%s*", args.CustomImage)
	}
	prefix := util.If(len(args.Prefix) > 0, args.Prefix, "main")
	r := rhelAIRequest{
		mCtx:             mCtx,
		prefix:           &prefix,
		amiName:          &amiName,
		arch:             &args.Arch,
		timeout:          &args.Timeout,
		serviceEndpoints: args.ServiceEndpoints,
		vpcID:            args.VpcID,
		diskSize:         args.ComputeRequest.DiskSize,
		model:            &args.Model,
		hfToken:          &args.HFToken,
		apiKey:           &args.APIKey,
		autoStart:        args.AutoStart,
		vllmExtraArgs:    &args.VLLMExtraArgs,
		exposePorts:      args.ExposePorts}
	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,
			VpcID:                 args.VpcID,
		})
	if err != nil {
		return err
	}
	if err = checkAMIExists(mCtx.Context(), &amiName, r.allocationData.Region, &amiArch); err != nil {
		return err
	}
	return r.createMachine()
}

// Will destroy resources related to machine
func Destroy(mCtxArgs *mc.ContextArgs) error {
	logging.Debug("Run rhel destroy")
	// Create mapt Context
	mCtx, err := mc.Init(mCtxArgs, aws.Provider())
	if err != nil {
		return err
	}

	if err := aws.DestroyStack(
		mCtx,
		aws.DestroyStackRequest{
			Stackname: stackName,
		}); err != nil {
		return err
	}
	if spot.Exist(mCtx) {
		if err := spot.Destroy(mCtx); err != nil {
			return err
		}
	}

	// Cleanup S3 state after all stacks have been destroyed
	return aws.CleanupState(mCtx)
}

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) {
	acc := strings.ToLower(strings.TrimSpace(accelerator))
	switch acc {
	case "cuda", "rocm":
	default:
		return nil, fmt.Errorf("unsupported accelerator %q (expected: cuda or rocm)", accelerator)
	}
	nameFilter := fmt.Sprintf("rhel-ai-%s-aws-*", acc)
	region := listVersionsRegion
	images, err := data.ListAMIs(ctx, data.ImageRequest{
		Name:   &nameFilter,
		Arch:   &amiArch,
		Owner:  &amiOwner,
		Region: &region,
	})
	if err != nil {
		return nil, fmt.Errorf("listing RHEL AI AMIs for accelerator %q: %w", acc, err)
	}
	prefix := fmt.Sprintf("rhel-ai-%s-aws-", acc)
	seen := make(map[string]struct{})
	for _, img := range images {
		if img.Name == nil {
			continue
		}
		raw := strings.TrimPrefix(*img.Name, prefix)
		// AMI names may have trailing qualifiers (e.g. "-x86_64-..."); take only the version part
		if idx := strings.Index(raw, "-x86_64"); idx > 0 {
			raw = raw[:idx]
		}
		if idx := strings.Index(raw, "-arm64"); idx > 0 {
			raw = raw[:idx]
		}
		if len(raw) > 0 {
			// Normalize underscores to dashes (e.g. "3.4.0_ea.2" → "3.4.0-ea.2")
			version := strings.ReplaceAll(raw, "_", "-")
			seen[version] = 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),
		ProjectName: r.mCtx.ProjectName(),
		BackedURL:   r.mCtx.BackedURL(),
		ProviderCredentials: aws.GetClouProviderCredentials(
			map[string]string{
				awsConstants.CONFIG_AWS_REGION: *r.allocationData.Region}),
		DeployFunc: r.deploy,
	}

	sr, _ := manager.UpStack(r.mCtx, cs)
	return r.manageResults(sr)
}

// function wil all the logic to deploy resources required by windows
// * create AMI Copy if needed
// * networking
// * key
// * security group
// * compute
// * checks
func (r *rhelAIRequest) deploy(ctx *pulumi.Context) error {
	if err := r.validate(); err != nil {
		return err
	}
	// Get AMI
	ami, err := amiSVC.GetAMIByName(ctx,
		*r.amiName,
		[]string{amiOwner},
		map[string]string{
			"architecture": amiArch})
	if err != nil {
		return err
	}
	// Networking
	nw, err := network.Create(ctx, r.mCtx,
		&network.NetworkArgs{
			Prefix:             *r.prefix,
			ID:                 awsRHELDedicatedID,
			Region:             *r.allocationData.Region,
			AZ:                 *r.allocationData.AZ,
			CreateLoadBalancer: r.allocationData.SpotPrice != nil,
			ServiceEndpoints:   r.serviceEndpoints,
			VpcID:              r.vpcID,
		})
	if err != nil {
		return err
	}
	// Create Keypair
	kpr := keypair.KeyPairRequest{
		Name: resourcesUtil.GetResourceName(
			*r.prefix, awsRHELDedicatedID, "pk")}
	keyResources, err := kpr.Create(ctx, r.mCtx)
	if err != nil {
		return err
	}
	ctx.Export(fmt.Sprintf("%s-%s", *r.prefix, outputUserPrivateKey),
		keyResources.PrivateKey.PrivateKeyPem)
	// Security groups
	securityGroups, err := r.securityGroups(ctx, r.mCtx, nw.Vpc, nw.IsPublic)
	if err != nil {
		return err
	}
	effectiveDiskSize := diskSize
	if r.diskSize != nil {
		effectiveDiskSize = *r.diskSize
	}
	cr := compute.ComputeRequest{
		MCtx:           r.mCtx,
		Prefix:         *r.prefix,
		ID:             awsRHELDedicatedID,
		VPC:            nw.Vpc,
		Subnet:         nw.Subnet,
		AMI:            ami,
		KeyResources:   keyResources,
		SecurityGroups: securityGroups,
		InstaceTypes:   r.allocationData.InstanceTypes,
		DiskSize:       &effectiveDiskSize,
		LB:             nw.LoadBalancer,
		Eip:            nw.Eip,
		LBTargetGroups: r.lbTargetGroups()}
	if r.allocationData.SpotPrice != nil && nw.IsPublic {
		cr.Spot = true
		cr.SpotPrice = *r.allocationData.SpotPrice
	}
	c, err := cr.NewCompute(ctx)
	if err != nil {
		return err
	}
	ctx.Export(fmt.Sprintf("%s-%s", *r.prefix, outputUsername),
		pulumi.String(amiUserDefault))
	ctx.Export(fmt.Sprintf("%s-%s", *r.prefix, outputHost),
		c.GetHostDnsName(true))
	if len(*r.timeout) > 0 {
		if err = serverless.OneTimeDelayedTask(ctx, r.mCtx,
			*r.allocationData.Region, *r.prefix,
			awsRHELDedicatedID,
			fmt.Sprintf("aws %s destroy --project-name %s --backed-url %s --serverless",
				"rhel",
				r.mCtx.ProjectName(),
				r.mCtx.BackedURL()),
			*r.timeout); err != nil {
			return err
		}
	}
	if !nw.IsPublic {
		if r.autoStart {
			return fmt.Errorf("--auto-start requires SSH access to configure RHAIIS; private subnet deployments do not support --auto-start")
		}
		return nil
	}
	if !r.autoStart {
		return c.Readiness(ctx, command.CommandPing, *r.prefix, awsRHELDedicatedID,
			keyResources.PrivateKey, amiUserDefault, nil, c.Dependencies)
	}
	readinessCmd, err := c.RunCommand(ctx,
		command.CommandPing,
		compute.LoggingCmdStd,
		fmt.Sprintf("%s-readiness", *r.prefix), awsRHELDedicatedID,
		keyResources.PrivateKey, amiUserDefault,
		nil, c.Dependencies)
	if err != nil {
		return err
	}
	_, err = c.RunCommand(ctx,
		r.rhaiisSetupScript(),
		compute.NoLoggingCmdStd,
		fmt.Sprintf("%s-rhaiis-setup", *r.prefix), awsRHELDedicatedID,
		keyResources.PrivateKey, amiUserDefault,
		nil, []pulumi.Resource{readinessCmd})
	return err
}

// Write exported values in context to files o a selected target folder
func (r *rhelAIRequest) manageResults(stackResult auto.UpResult) error {
	results := map[string]string{
		fmt.Sprintf("%s-%s", *r.prefix, outputUsername):       "username",
		fmt.Sprintf("%s-%s", *r.prefix, outputUserPrivateKey): "id_rsa",
		fmt.Sprintf("%s-%s", *r.prefix, outputHost):           "host",
	}
	return output.Write(stackResult, r.mCtx.GetResultsOutputPath(), results)
}

func (r *rhelAIRequest) securityGroups(ctx *pulumi.Context, mCtx *mc.Context,
	vpc *ec2.Vpc, public bool) (pulumi.StringArray, error) {
	var ingressRules []securityGroup.IngressRules
	if public {
		sshIngressRule := securityGroup.SSH_TCP
		sshIngressRule.CidrBlocks = infra.NETWORKING_CIDR_ANY_IPV4
		ingressRules = []securityGroup.IngressRules{sshIngressRule}
		for _, port := range r.exposePorts {
			ingressRules = append(ingressRules, securityGroup.IngressRules{
				Description: fmt.Sprintf("port-%d", port),
				FromPort:    port,
				ToPort:      port,
				Protocol:    "tcp",
				CidrBlocks:  infra.NETWORKING_CIDR_ANY_IPV4,
			})
		}
	}
	// Create SG with ingress rules
	sg, err := securityGroup.SGRequest{
		Name:        resourcesUtil.GetResourceName(*r.prefix, awsRHELDedicatedID, "sg"),
		VPC:         vpc,
		Description: fmt.Sprintf("sg for %s", awsRHELDedicatedID),
		IngressRules: ingressRules,
	}.Create(ctx, mCtx)
	if err != nil {
		return nil, err
	}
	// Convert to an array of IDs
	sgs := util.ArrayConvert([]*ec2.SecurityGroup{sg.SG},
		func(sg *ec2.SecurityGroup) pulumi.StringInput {
			return sg.ID()
		})
	return pulumi.StringArray(sgs[:]), nil
}

func (r *rhelAIRequest) lbTargetGroups() []int {
	return append([]int{22}, r.exposePorts...)
}

func (r *rhelAIRequest) rhaiisSetupScript() string {
	confDir := "/etc/containers/systemd/rhaii.container.d"
	script := fmt.Sprintf(
		"sudo cp %s/install.conf.example %s/install.conf",
		confDir, confDir)
	if len(*r.hfToken) > 0 {
		script += fmt.Sprintf(
			" && sudo sed -i 's|HUGGING_FACE_HUB_TOKEN=.*|HUGGING_FACE_HUB_TOKEN=%s|' %s/install.conf",
			*r.hfToken, confDir)
	}
	if len(*r.model) > 0 {
		script += fmt.Sprintf(
			` && sudo sed -i 's|--model .*|--model %s \\|' %s/install.conf`,
			*r.model, confDir)
	}
	script += fmt.Sprintf(
		` && GPU_COUNT=$(nvidia-smi -L 2>/dev/null | wc -l) && [ "$GPU_COUNT" -gt 0 ] && sudo sed -i "s|--tensor-parallel-size 1|--tensor-parallel-size $GPU_COUNT|" %s/install.conf`,
		confDir)
	if len(*r.vllmExtraArgs) > 0 {
		extraArgs := *r.vllmExtraArgs
		if strings.Contains(extraArgs, "--max-model-len") {
			script += fmt.Sprintf(
				` && sudo sed -i 's|--max-model-len [0-9]*|%s|' %s/install.conf`,
				extraArgs, confDir)
		} else {
			script += fmt.Sprintf(
				` && sudo sed -i 's|--max-model-len 4096|--max-model-len 4096 \\\n     %s|' %s/install.conf`,
				extraArgs, confDir)
		}
	}
	if len(*r.apiKey) > 0 {
		script += fmt.Sprintf(
			" && sudo sed -i '/\\[Install\\]/i Environment=VLLM_API_KEY=%s' %s/install.conf",
			*r.apiKey, confDir)
	}
	script += " && sudo systemctl daemon-reload && sudo systemctl start rhaii"
	return script
}

func checkAMIExists(ctx context.Context, amiName, region, arch *string) error {
	isAMIOffered, _, err := data.IsAMIOffered(
		ctx,
		data.ImageRequest{
			Name:   amiName,
			Arch:   arch,
			Region: region,
			Owner:  &amiOwner})
	if err != nil {
		return err
	}
	if !isAMIOffered {
		return fmt.Errorf("AMI %s could not be found in region: %s", *amiName, *region)
	}
	return nil
}