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')