github/kubernetes
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow passing timeout on TPU API operations via context

Browse Source
This commit is contained in:
Yang Guo
2018-08-13 15:36:25 -07:00
parent facc84752b
commit c776cf4296
1 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
pkg/cloudprovider/providers/gce/gce_tpu.go
View File

@@ -65,7 +65,7 @@ func (gce *GCECloud) CreateTPU(ctx context.Context, name, zone string, node *tpu
}
glog.V(2).Infof("Creating Cloud TPU %q in zone %q with operation %q", name, zone, op.Name)
op, err = gce.waitForTPUOp(30*time.Second, 10*time.Minute, op)
op, err = gce.waitForTPUOp(ctx, op)
if err != nil {
return nil, err
}
@@ -98,7 +98,7 @@ func (gce *GCECloud) DeleteTPU(ctx context.Context, name, zone string) error {
}
glog.V(2).Infof("Deleting Cloud TPU %q in zone %q with operation %q", name, zone, op.Name)
op, err = gce.waitForTPUOp(30*time.Second, 10*time.Minute, op)
op, err = gce.waitForTPUOp(ctx, op)
if err != nil {
return err
}
@@ -133,10 +133,18 @@ func (gce *GCECloud) ListTPUs(ctx context.Context, zone string) ([]*tpuapi.Node,
return response.Nodes, mc.Observe(nil)
}
// waitForTPUOp checks whether the op is done every interval before the timeout
// occurs.
func (gce *GCECloud) waitForTPUOp(interval, timeout time.Duration, op *tpuapi.Operation) (*tpuapi.Operation, error) {
if err := wait.PollImmediate(interval, timeout, func() (bool, error) {
// 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) {
if err := wait.PollInfinite(30*time.Second, func() (bool, error) {
// Check if context has been cancelled.
select {
case <-ctx.Done():
glog.V(3).Infof("Context for operation %q has been cancelled: %s", op.Name, ctx.Err())
return true, ctx.Err()
default:
}
glog.V(3).Infof("Waiting for operation %q to complete...", op.Name)
start := time.Now()