Switch GPU selection to model instead of counts (#2807)

* Switch GPU selection to model instead of counts

* update gpu model names

* rename to --vm-gpu-kind

Author: dangra

api/machine_types.go

@@ -25,6 +25,7 @@ const (
 	MachineStateStopped                        = "stopped"
 	MachineStateCreated                        = "created"
 	DefaultVMSize                              = "shared-cpu-1x"
+	DefaultGPUVMSize                           = "performance-8x"
 )
 
 type Machine struct {
@@ -321,7 +322,7 @@ type MachineGuest struct {
 	CPUKind  string `json:"cpu_kind,omitempty"`
 	CPUs     int    `json:"cpus,omitempty"`
 	MemoryMB int    `json:"memory_mb,omitempty"`
-	GPUs     int    `json:"gpus,omitempty"`
+	GPUKind  string `json:"gpu_kind,omitempty"`
 
 	KernelArgs []string `json:"kernel_args,omitempty"`
 }

internal/command/console/console.go

@@ -240,7 +240,7 @@ func determineEphemeralConsoleMachineGuest(ctx context.Context) (*api.MachineGue
 		}
 	}
 
-	if cpuKind := flag.GetString(ctx, "vm-cpukind"); cpuKind != "" {
+	if cpuKind := flag.GetString(ctx, "vm-cpu-kind"); cpuKind != "" {
 		desiredGuest.CPUKind = cpuKind
 	}
 

internal/flag/machines.go

@@ -3,45 +3,61 @@ package flag
 import (
 	"context"
 	"fmt"
+	"slices"
+	"strings"
 
 	"github.com/superfly/flyctl/api"
 )
 
+var validGPUKinds = []string{"a100-pcie-40gb", "a100-sxm4-80gb"}
+
 // Returns a MachineGuest based on the flags provided overwriting a default VM
 func GetMachineGuest(ctx context.Context, guest *api.MachineGuest) (*api.MachineGuest, error) {
+	defaultVMSize := api.DefaultVMSize
+	if IsSpecified(ctx, "vm-gpu-kind") {
+		defaultVMSize = api.DefaultGPUVMSize
+	}
+
 	if guest == nil {
 		guest = &api.MachineGuest{}
-		guest.SetSize(api.DefaultVMSize)
+		guest.SetSize(defaultVMSize)
 	}
 
 	if IsSpecified(ctx, "vm-size") {
 		if err := guest.SetSize(GetString(ctx, "vm-size")); err != nil {
 			return nil, err
 		}
 	}
+
 	if IsSpecified(ctx, "vm-cpus") {
 		guest.CPUs = GetInt(ctx, "vm-cpus")
 		if guest.CPUs == 0 {
-			return nil, fmt.Errorf("cannot have zero cpus")
+			return nil, fmt.Errorf("--vm-cpus cannot be zero")
 		}
 	}
 
 	if IsSpecified(ctx, "vm-memory") {
 		guest.MemoryMB = GetInt(ctx, "vm-memory")
 		if guest.MemoryMB == 0 {
-			return nil, fmt.Errorf("memory cannot be zero")
+			return nil, fmt.Errorf("--vm-memory cannot be zero")
 		}
 	}
 
-	if IsSpecified(ctx, "vm-cpukind") {
-		guest.CPUKind = GetString(ctx, "vm-cpukind")
+	if IsSpecified(ctx, "vm-cpu-kind") {
+		guest.CPUKind = GetString(ctx, "vm-cpu-kind")
 		if k := guest.CPUKind; k != "shared" && k != "performance" {
-			return nil, fmt.Errorf("cpukind must be set to 'shared' or 'performance'")
+			return nil, fmt.Errorf("--vm-cpu-kind must be set to 'shared' or 'performance'")
 		}
 	}
-	if IsSpecified(ctx, "vm-gpus") {
-		guest.GPUs = GetInt(ctx, "vm-gpus")
+
+	if IsSpecified(ctx, "vm-gpu-kind") {
+		m := GetString(ctx, "vm-gpu-kind")
+		if !slices.Contains(validGPUKinds, m) {
+			return nil, fmt.Errorf("--vm-gpu-kind must be set to one of: %v", strings.Join(validGPUKinds, ", "))
+		}
+		guest.GPUKind = m
 	}
+
 	return guest, nil
 }
 
@@ -57,16 +73,18 @@ var VMSizeFlags = Set{
 		Aliases:     []string{"cpus"},
 	},
 	String{
-		Name:        "vm-cpukind",
+		Name:        "vm-cpu-kind",
 		Description: "The kind of CPU to use ('shared' or 'performance')",
+		Aliases:     []string{"vm-cpukind"},
 	},
 	Int{
 		Name:        "vm-memory",
 		Description: "Memory (in megabytes) to attribute to the VM",
 		Aliases:     []string{"memory"},
 	},
-	Int{
-		Name:        "vm-gpus",
-		Description: "Number of GPUs",
+	String{
+		Name:        "vm-gpu-kind",
+		Description: fmt.Sprintf("If set, the GPU model to attach (%v)", strings.Join(validGPUKinds, ", ")),
+		Aliases:     []string{"vm-gpukind"},
 	},
 }

test/preflight/apps_v2_integration_test.go

@@ -487,7 +487,7 @@ func TestLaunchCpusMem(t *testing.T) {
 	f := testlib.NewTestEnvFromEnv(t)
 	appName := f.CreateRandomAppName()
 
-	f.Fly("launch --org %s --name %s --region %s --now --internal-port 80 --image nginx --auto-confirm --vm-cpus 4 --vm-memory 8192 --vm-cpukind performance", f.OrgSlug(), appName, f.PrimaryRegion())
+	f.Fly("launch --org %s --name %s --region %s --now --internal-port 80 --image nginx --auto-confirm --vm-cpus 4 --vm-memory 8192 --vm-cpu-kind performance", f.OrgSlug(), appName, f.PrimaryRegion())
 	machines := f.MachinesList(appName)
 	firstMachineGuest := machines[0].Config.Guest