From 8b6b81c7c982e9a8c0190efa0aca31627e7f1002 Mon Sep 17 00:00:00 2001 From: Yang Guo Date: Mon, 27 Aug 2018 16:56:18 -0700 Subject: [PATCH] GCE: Add ListLocations to Cloud TPU API --- pkg/cloudprovider/providers/gce/gce_tpu.go | 31 +++++++++++++++------- 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/pkg/cloudprovider/providers/gce/gce_tpu.go b/pkg/cloudprovider/providers/gce/gce_tpu.go index 01d3cce275f..0a78f62cb3a 100644 --- a/pkg/cloudprovider/providers/gce/gce_tpu.go +++ b/pkg/cloudprovider/providers/gce/gce_tpu.go @@ -38,16 +38,14 @@ func newTPUService(client *http.Client) (*tpuService, error) { return nil, err } return &tpuService{ - nodesService: tpuapi.NewProjectsLocationsNodesService(s), - operationsService: tpuapi.NewProjectsLocationsOperationsService(s), + projects: tpuapi.NewProjectsService(s), }, nil } // tpuService encapsulates the TPU services on nodes and the operations on the // nodes. type tpuService struct { - nodesService *tpuapi.ProjectsLocationsNodesService - operationsService *tpuapi.ProjectsLocationsOperationsService + projects *tpuapi.ProjectsService } // CreateTPU creates the Cloud TPU node with the specified name in the @@ -59,7 +57,7 @@ func (gce *GCECloud) CreateTPU(ctx context.Context, name, zone string, node *tpu var op *tpuapi.Operation parent := getTPUParentName(gce.projectID, zone) - op, err = gce.tpuService.nodesService.Create(parent, node).NodeId(name).Do() + op, err = gce.tpuService.projects.Locations.Nodes.Create(parent, node).NodeId(name).Do() if err != nil { return nil, err } @@ -92,7 +90,7 @@ func (gce *GCECloud) DeleteTPU(ctx context.Context, name, zone string) error { var op *tpuapi.Operation name = getTPUName(gce.projectID, zone, name) - op, err = gce.tpuService.nodesService.Delete(name).Do() + op, err = gce.tpuService.projects.Locations.Nodes.Delete(name).Do() if err != nil { return err } @@ -114,7 +112,7 @@ func (gce *GCECloud) GetTPU(ctx context.Context, name, zone string) (*tpuapi.Nod mc := newTPUMetricContext("get", zone) name = getTPUName(gce.projectID, zone, name) - node, err := gce.tpuService.nodesService.Get(name).Do() + node, err := gce.tpuService.projects.Locations.Nodes.Get(name).Do() if err != nil { return nil, mc.Observe(err) } @@ -126,13 +124,24 @@ func (gce *GCECloud) ListTPUs(ctx context.Context, zone string) ([]*tpuapi.Node, mc := newTPUMetricContext("list", zone) parent := getTPUParentName(gce.projectID, zone) - response, err := gce.tpuService.nodesService.List(parent).Do() + response, err := gce.tpuService.projects.Locations.Nodes.List(parent).Do() if err != nil { return nil, mc.Observe(err) } return response.Nodes, mc.Observe(nil) } +// ListLocations returns the zones where Cloud TPUs are available. +func (gce *GCECloud) ListLocations(ctx context.Context) ([]*tpuapi.Location, error) { + mc := newTPUMetricContext("list_locations", "") + parent := getTPUProjectURL(gce.projectID) + response, err := gce.tpuService.projects.Locations.List(parent).Do() + if err != nil { + return nil, mc.Observe(err) + } + return response.Locations, mc.Observe(nil) +} + // waitForTPUOp checks whether the op is done every 30 seconds before the ctx // is cancelled. func (gce *GCECloud) waitForTPUOp(ctx context.Context, op *tpuapi.Operation) (*tpuapi.Operation, error) { @@ -155,7 +164,7 @@ func (gce *GCECloud) waitForTPUOp(ctx context.Context, op *tpuapi.Operation) (*t } var err error - op, err = gce.tpuService.operationsService.Get(op.Name).Do() + op, err = gce.tpuService.projects.Locations.Operations.Get(op.Name).Do() if err != nil { return true, err } @@ -188,6 +197,10 @@ func getErrorFromTPUOp(op *tpuapi.Operation) error { return nil } +func getTPUProjectURL(project string) string { + return fmt.Sprintf("projects/%s", project) +} + func getTPUParentName(project, zone string) string { return fmt.Sprintf("projects/%s/locations/%s", project, zone) }