def Describe(tpu_node, zone=None):
"""Invoke TPU Get API."""
zone = zone or properties.VALUES.compute.zone.GetOrFail
tpu_api_client = api_util.TpusClient('v1alpha1')
node_ref = resources.REGISTRY.Parse(
tpu_node,
params={
'projectsId': properties.VALUES.core.project.GetOrFail,
'locationsId': zone},
collection=TPU_NODE_COLLECTION)
return tpu_api_client.Get(node_ref)
def List(page_size, limit, zone=None):
"""Invoke TPU List API."""
zone = zone or properties.VALUES.compute.zone.GetOrFail()
tpu_api_client = api_util.TpusClient('v1alpha1')
location_ref = resources.REGISTRY.Parse(
zone,
params={
'projectsId': properties.VALUES.core.project.GetOrFail,
'locationsId': zone},
collection=TPU_LOCATION_COLLECTION)
return tpu_api_client.List(location_ref, page_size, limit)
def WaitForOperation(operation, zone):
"""Wait for the specified tpu operation."""
wait_message = 'Waiting for [{0}] to finish'.format(operation.name)
tpu_api_client = api_util.TpusClient('v1alpha1')
poller = TpuOperationsPoller(tpu_api_client)
operation_ref = resources.REGISTRY.Parse(
operation.name,
params={
'projectsId': properties.VALUES.core.project.GetOrFail,
'locationsId': zone
},
collection=TPU_OPERATION_COLLECTION)
return waiter.WaitFor(poller, operation_ref, wait_message)
def Create(name,
cidr_range,
description=None,
network='default',
accelerator_type=None,
version=None,
zone=None):
"""Invoke TPU Create API and return created resource."""
tpu_api_client = api_util.TpusClient('v1alpha1')
zone = zone or properties.VALUES.compute.zone.GetOrFail()
node_msg = tpu_api_client.messages.Node(cidrBlock=cidr_range,
network=network,
acceleratorType=accelerator_type,
tensorflowVersion=version,
description=description)
parent_ref = resources.REGISTRY.Parse(
zone,
params={'projectsId': properties.VALUES.core.project.GetOrFail},
collection=TPU_LOCATION_COLLECTION)
return WaitForOperation(tpu_api_client.Create(node_msg, parent_ref, name),
zone)
def Reimage(node_ref, version, zone):
"""Invoke TPU Reimage API and wait for operation result."""
tpu_api_client = api_util.TpusClient('v1alpha1')
return WaitForOperation(tpu_api_client.Reimage(node_ref, version), zone)
