def testGetNumShards(self):
with self._PatchCriticalObjects(
stdout='{"networkEndpoints": [{"ipAddress":'
' "10.199.12.2", "port": 8470}]}'):
tpu = gcp_tpu.GcpTpu(self.mock_tpu_spec)
num_shards = tpu.GetNumShards()
self.assertEqual(num_shards, 8)
def testDelete(self):
with self._PatchCriticalObjects() as issue_command:
tpu = gcp_tpu.GcpTpu(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 compute tpus delete pkb-tpu-123'))
self.assertIn('--project fakeproject', command_string)
self.assertIn('--zone us-central1-a', command_string)
def testGetMasterGrpcAddress(self):
with self._PatchCriticalObjects(stdout="""{
"networkEndpoints": [{
"ipAddress": "10.199.12.2",
"port": 8470
}]
}
"""):
tpu = gcp_tpu.GcpTpu(self.mock_tpu_spec)
ip_address = tpu.GetMasterGrpcAddress()
self.assertEqual(ip_address, 'grpc://10.199.12.2:8470')
def testStockout(self):
stderr = """Create request issued for: [pkb-tpu-train-9baf32202]
Waiting for operation [projects/artemis-prod/locations/us-central1-b/operations/operation-1567697651843-591d00da740fa-ed64d57f-8a2533cb] to complete failed.
ERROR: (gcloud.compute.tpus.create) {
"code": 8,
"message": "There is no more capacity in the zone \"us-central1-b\"; you can try in another zone where Cloud TPU Nodes are offered (see https://cloud.google.com/tpu/docs/regions) [EID: 0xf3bb52b78a15cd16]"
}"""
with self._PatchCriticalObjects(stderr=stderr, return_code=1):
with self.assertRaises(
errors.Benchmarks.InsufficientCapacityCloudFailure):
tpu = gcp_tpu.GcpTpu(self.mock_tpu_spec)
tpu._Create()
def testCreate(self):
with self._PatchCriticalObjects() as issue_command:
tpu = gcp_tpu.GcpTpu(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)
self.assertIn('--preemptible', command_string)
def CreateTpuFromSpec(self, spec_dict):
mock_tpu_spec = mock.Mock(spec=benchmark_config_spec._TpuGroupSpec)
mock_tpu_spec.configure_mock(**spec_dict)
tpu_class = gcp_tpu.GcpTpu(mock_tpu_spec)
return tpu_class
def testGetName(self):
with self._PatchCriticalObjects():
tpu = gcp_tpu.GcpTpu(self.mock_tpu_spec)
name = tpu.GetName()
self.assertEqual(name, 'pkb-tpu-123')
