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