diff --git a/api/cloud.go b/api/cloud.go index e73b2b7..95b393f 100644 --- a/api/cloud.go +++ b/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 +} diff --git a/cmd/cloud/getCloud.go b/cmd/cloud/getCloud.go index 831c9c4..b996fbf 100644 --- a/cmd/cloud/getCloud.go +++ b/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") }