Add cloud option to show datacenter for the gpus

api/cloud.go

@@ -68,3 +68,69 @@ func GetCloud(in *GetCloudInput) (gpuTypes []interface{}, err error) {
 	}
 	return
 }
+
+// GetDCsByGPU returns a mapping of GPU ID to dataCenterIDs.
+// If there is an API failure, the map will simply be empty.
+func GetDCsByGPU() (mapping map[string][]string) {
+	mapping = make(map[string][]string)
+
+	input := Input{
+		Query: `
+		  query GetDcGPUs {
+			dataCenters {
+			  id
+			  gpuAvailability {
+			    gpuTypeId
+			  }
+			}
+		  }
+		`,
+	}
+
+	res, err := Query(input)
+	if err != nil {
+		return
+	}
+	defer res.Body.Close()
+	rawData, err := io.ReadAll(res.Body)
+	if err != nil {
+		return
+	}
+	if res.StatusCode != 200 {
+		return
+	}
+
+	type response struct {
+		Data struct {
+			DataCenters []struct {
+				ID              string `json:"id"`
+				GpuAvailability []struct {
+					GpuTypeID string `json:"gpuTypeId"`
+				} `json:"gpuAvailability"`
+			} `json:"dataCenters"`
+			Errors []struct {
+				Message string `json:"message"`
+			} `json:"errors"`
+		} `json:"data"`
+	}
+
+	var r response
+	if err = json.Unmarshal(rawData, &r); err != nil {
+		return
+	}
+
+	if r.Data.Errors != nil {
+		return
+	}
+
+	for _, dc := range r.Data.DataCenters {
+		for _, ga := range dc.GpuAvailability {
+			gpu := ga.GpuTypeID
+			if _, ok := mapping[gpu]; !ok {
+				mapping[gpu] = []string{}
+			}
+			mapping[gpu] = append(mapping[gpu], dc.ID)
+		}
+	}
+	return mapping
+}

cmd/cloud/getCloud.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"strconv"
+	"strings"
 
 	"github.com/runpod/runpodctl/api"
 	"github.com/runpod/runpodctl/format"
@@ -18,6 +19,7 @@ var (
 	memory    int
 	vcpu      int
 	secure    bool
+	showDC    bool
 )
 
 var GetCloudCmd = &cobra.Command{
@@ -49,6 +51,11 @@ var GetCloudCmd = &cobra.Command{
 		gpuTypes, err := api.GetCloud(input)
 		cobra.CheckErr(err)
 
+		var dcs dcDecorator
+		if showDC {
+			dcs = dcDecorator{api.GetDCsByGPU()}
+		}
+
 		data := [][]string{}
 		for _, gpu := range gpuTypes {
 			gpuType, ok := gpu.(map[string]interface{})
@@ -76,10 +83,12 @@ var GetCloudCmd = &cobra.Command{
 				spotPriceString,
 				onDemandPriceString,
 			}
+			row = dcs.decorateRow(row, kv["gpuTypeId"].(string))
 			data = append(data, row)
 		}
 
 		header := []string{"GPU Type", "Mem GB", "vCPU", "Spot $/HR", "OnDemand $/HR"}
+		header = dcs.decorateHeader(header, "Data centers")
 		tb := tablewriter.NewWriter(os.Stdout)
 		tb.SetHeader(header)
 		tb.AppendBulk(data)
@@ -88,10 +97,34 @@ var GetCloudCmd = &cobra.Command{
 	},
 }
 
+type dcDecorator struct {
+	dcsByGPU map[string][]string
+}
+
+func (d *dcDecorator) decorateRow(fields []string, gpuID string) []string {
+	if d.dcsByGPU == nil {
+		return fields
+	}
+	dcs, ok := d.dcsByGPU[gpuID]
+	if !ok {
+		return fields
+	}
+
+	return append(fields, strings.Join(dcs, ","))
+}
+
+func (d *dcDecorator) decorateHeader(headers []string, columnName string) []string {
+	if d.dcsByGPU == nil {
+		return headers
+	}
+	return append(headers, columnName)
+}
+
 func init() {
 	GetCloudCmd.Flags().BoolVarP(&community, "community", "c", false, "show listings from community cloud only")
 	GetCloudCmd.Flags().IntVar(&disk, "disk", 0, "minimum disk size in GB you need")
 	GetCloudCmd.Flags().IntVar(&memory, "mem", 0, "minimum sys memory size in GB you need")
 	GetCloudCmd.Flags().IntVar(&vcpu, "vcpu", 0, "minimum vCPUs you need")
 	GetCloudCmd.Flags().BoolVarP(&secure, "secure", "s", false, "show listings from secure cloud only")
+	GetCloudCmd.Flags().BoolVarP(&showDC, "datacenter", "d", false, "show which datacenters provide which gpu types")
 }