From fc3c8fba031fb5680873d3835009e055b5e6fa9d Mon Sep 17 00:00:00 2001 From: Jonas Tingeborn Date: Thu, 27 Feb 2025 11:00:49 +0100 Subject: [PATCH 1/2] Add cloud option to show datacenter for the gpus --- api/cloud.go | 66 +++++++++++++++++++++++++++++++++++++++++++ cmd/cloud/getCloud.go | 23 +++++++++++++++ 2 files changed, 89 insertions(+) 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..95c771e 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 dcsByGPU map[string][]string + if showDC { + dcsByGPU = api.GetDCsByGPU() + } + data := [][]string{} for _, gpu := range gpuTypes { gpuType, ok := gpu.(map[string]interface{}) @@ -76,10 +83,16 @@ var GetCloudCmd = &cobra.Command{ spotPriceString, onDemandPriceString, } + if dcsByGPU != nil { + row = decorateDataCenters(row, kv["gpuTypeId"].(string), dcsByGPU) + } data = append(data, row) } header := []string{"GPU Type", "Mem GB", "vCPU", "Spot $/HR", "OnDemand $/HR"} + if dcsByGPU != nil { + header = append(header, "Data centers") + } tb := tablewriter.NewWriter(os.Stdout) tb.SetHeader(header) tb.AppendBulk(data) @@ -88,10 +101,20 @@ var GetCloudCmd = &cobra.Command{ }, } +func decorateDataCenters(fields []string, gpuID string, dcsByGPU map[string][]string) []string { + dcs, ok := dcsByGPU[gpuID] + if !ok { + return fields + } + + return append(fields, strings.Join(dcs, ",")) +} + 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") } From cbe18bbf816f48e59311d068ec541fc8df1e405f Mon Sep 17 00:00:00 2001 From: Jonas Tingeborn Date: Thu, 27 Feb 2025 15:56:30 +0100 Subject: [PATCH 2/2] Factor out conditional cloud DC column generation Reduces branching in the main flow. --- cmd/cloud/getCloud.go | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/cmd/cloud/getCloud.go b/cmd/cloud/getCloud.go index 95c771e..b996fbf 100644 --- a/cmd/cloud/getCloud.go +++ b/cmd/cloud/getCloud.go @@ -51,9 +51,9 @@ var GetCloudCmd = &cobra.Command{ gpuTypes, err := api.GetCloud(input) cobra.CheckErr(err) - var dcsByGPU map[string][]string + var dcs dcDecorator if showDC { - dcsByGPU = api.GetDCsByGPU() + dcs = dcDecorator{api.GetDCsByGPU()} } data := [][]string{} @@ -83,16 +83,12 @@ var GetCloudCmd = &cobra.Command{ spotPriceString, onDemandPriceString, } - if dcsByGPU != nil { - row = decorateDataCenters(row, kv["gpuTypeId"].(string), dcsByGPU) - } + row = dcs.decorateRow(row, kv["gpuTypeId"].(string)) data = append(data, row) } header := []string{"GPU Type", "Mem GB", "vCPU", "Spot $/HR", "OnDemand $/HR"} - if dcsByGPU != nil { - header = append(header, "Data centers") - } + header = dcs.decorateHeader(header, "Data centers") tb := tablewriter.NewWriter(os.Stdout) tb.SetHeader(header) tb.AppendBulk(data) @@ -101,8 +97,15 @@ var GetCloudCmd = &cobra.Command{ }, } -func decorateDataCenters(fields []string, gpuID string, dcsByGPU map[string][]string) []string { - dcs, ok := dcsByGPU[gpuID] +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 } @@ -110,6 +113,13 @@ func decorateDataCenters(fields []string, gpuID string, dcsByGPU map[string][]st 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")