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)