def testDelete(self):
with self._PatchCriticalObjects() as issue_command:
tpu = gcp_cloud_tpu.GcpCloudTpu(self.mock_tpu_spec)
tpu._Delete()
self.assertEqual(issue_command.call_count, 1)
command_string = ' '.join(issue_command.call_args[0][0])
self.assertTrue(
command_string.startswith(
'gcloud alpha compute tpus delete pkb-tpu-123'))
self.assertIn('--project fakeproject', command_string)
self.assertIn('--zone us-central1-a', command_string)
def testCreate(self):
with self._PatchCriticalObjects() as issue_command:
tpu = gcp_cloud_tpu.GcpCloudTpu(self.mock_tpu_spec)
tpu._Create()
self.assertEqual(issue_command.call_count, 1)
command_string = ' '.join(issue_command.call_args[0][0])
self.assertTrue(
command_string.startswith(
'gcloud compute tpus create pkb-tpu-123'), command_string)
self.assertIn('--project fakeproject', command_string)
self.assertIn('--range 192.168.0.0/29', command_string)
self.assertIn('--accelerator-type tpu-v2', command_string)
self.assertIn('--description MyTFNode', command_string)
self.assertIn('--network default', command_string)
self.assertIn('--version nightly', command_string)
self.assertIn('--zone us-central1-a', command_string)
def testGetCloudTpuIp(self):
with self._PatchCriticalObjects(stdout=json.dumps({
'@type':
'type.googleapis.com/google.cloud.tpu.v1alpha1.Node',
'acceleratorType':
'zones/us-central1-a/acceleratorTypes/tpu-v2',
'cidrBlock':
'10.240.0.0/29',
'createTime':
'2017-10-02T20:48:54.421627Z',
'greenVmInstanceId':
'4423519461898655909',
'greenVmSelflink': (
'https://www.googleapis.com/compute/alpha/projects/'
'g9250ab742a5cc428-tp/zones/us-central1-a/instances'
'/n-74e6bf72-w-0'),
'ipAddress':
'196.168.0.2',
'machineType':
'zones/us-central1-a/machineTypes/custom-64-425984',
'name': ('projects/fakeproject/locations/us-central1-a/nodes/'
'pkb-tpu-123'),
'network':
'global/networks/default',
'port':
'8470',
'serviceAccount':
'*****@*****.**',
'tensorflowVersion':
'nightly'
})) as issue_command:
tpu = gcp_cloud_tpu.GcpCloudTpu(self.mock_tpu_spec)
tpu_ip = tpu.GetCloudTpuIp()
self.assertEqual(issue_command.call_count, 1)
command_string = ' '.join(issue_command.call_args[0][0])
self.assertTrue(
command_string.startswith(
'gcloud alpha compute tpus describe pkb-tpu-123'))
self.assertIn('--project fakeproject', command_string)
self.assertIn('--zone us-central1-a', command_string)
self.assertEqual(tpu_ip, '196.168.0.2')
def CreateCloudTpuFromSpec(self, spec_dict):
mock_tpu_spec = mock.Mock(spec=benchmark_config_spec._CloudTpuSpec)
mock_tpu_spec.configure_mock(**spec_dict)
tpu_class = gcp_cloud_tpu.GcpCloudTpu(mock_tpu_spec)
return tpu_class
def testGetName(self):
with self._PatchCriticalObjects():
tpu = gcp_cloud_tpu.GcpCloudTpu(self.mock_tpu_spec)
name = tpu.GetName()
self.assertEqual(name, 'pkb-tpu-123')
