Added accelerator flag to cli and cvdr.toml

Signed-off-by: Johnny Yip <jcyip@google.com>
google · Jan 23, 2025 · cd9f950 · cd9f950
1 parent 7640c15
commit cd9f950
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go
@@ -84,6 +84,7 @@ const (
 	gcpMachineTypeFlag    = "gcp_machine_type"
 	gcpMinCPUPlatformFlag = "gcp_min_cpu_platform"
 	gcpBootDiskSizeGB     = "gcp_boot_disk_size_gb"
+	gcpAcceleratorFlag    = "gcp_accelerator"
 )
 
 const (
@@ -542,6 +543,7 @@ func hostCommand(opts *subCommandOpts) *cobra.Command {
 }
 
 func cvdCommands(opts *subCommandOpts) []*cobra.Command {
+	gcpAcceleratorFlagValues := []string{}
 	// Create command
 	createFlags := &CreateCVDFlags{
 		ServiceFlags:   opts.ServiceFlags,
@@ -553,6 +555,11 @@ func cvdCommands(opts *subCommandOpts) []*cobra.Command {
 		Short:   "Creates a CVD",
 		PreRunE: preRunE(createFlags, &opts.ServiceFlags.Service, &opts.InitialConfig),
 		RunE: func(c *cobra.Command, args []string) error {
+			configs, err := parseAcceleratorFlag(gcpAcceleratorFlagValues)
+			if err != nil {
+				return err
+			}
+			createFlags.GCP.AcceleratorConfigs = configs
 			return runCreateCVDCommand(c, args, createFlags, opts)
 		},
 	}
@@ -645,6 +652,8 @@ func cvdCommands(opts *subCommandOpts) []*cobra.Command {
 		create.Flags().StringVar(f.ValueRef, name, f.Default, f.Desc)
 		create.MarkFlagsMutuallyExclusive(hostFlag, name)
 	}
+	create.Flags().StringSliceVar(&gcpAcceleratorFlagValues, "host_" + gcpAcceleratorFlag,
+		opts.InitialConfig.DefaultService().Host.GCP.AcceleratorConfigs, acceleratorFlagDesc)
 	// List command
 	listFlags := &ListCVDsFlags{ServiceFlags: opts.ServiceFlags}
 	list := &cobra.Command{
@@ -1728,7 +1737,7 @@ type acceleratorConfig struct {
 	Type  string
 }
 
-// Values should follow the pattern: `type=[TYPE],count=[COUNT]`, i.e: `type=nvidia-tesla-p100,count=1`.
+// Values should follow the pattern: '"type=[TYPE],count=[COUNT]"', i.e: '"type=nvidia-tesla-p100,count=1"'.
 func parseAcceleratorFlag(values []string) ([]acceleratorConfig, error) {
 	singleValueParser := func(value string) (*acceleratorConfig, error) {
 		sanitized := strings.Join(strings.Fields(value), "")

diff --git a/pkg/cli/config.go b/pkg/cli/config.go
@@ -27,9 +27,10 @@ import (
 )
 
 type GCPHostConfig struct {
-	MachineType    string `json:"machine_type,omitempty"`
-	MinCPUPlatform string `json:"min_cpu_platform,omitempty"`
-	BootDiskSizeGB int64  `json:"boot_disk_size_gb,omitempty"`
+	MachineType        string   `json:"machine_type,omitempty"`
+	MinCPUPlatform     string   `json:"min_cpu_platform,omitempty"`
+	BootDiskSizeGB     int64    `json:"boot_disk_size_gb,omitempty"`
+	AcceleratorConfigs []string `json:"accelerator_configs,omitempty"`
 }
 
 type HostConfig struct {