def testHandlesByteStrings(self): self.assertEqual( client.Client(tpu='tpu_name', zone='zone', project='project')._full_name(), client.Client(tpu=b'tpu_name', zone=b'zone', project=b'project')._full_name(), )
def testRecoverableHBMOOMNoAPI(self): test_cases = [ ({ 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'state': 'READY', 'symptoms': [{ 'createTime': '2000-01-01T00:29:30.123456Z', 'symptomType': 'HBM_OUT_OF_MEMORY', 'details': 'The TPU HBM has run OOM at timestamp ' '2020-05-29T04:51:32.038721+00:00', 'workerId': '0' }] } }, True), ] for tpu_map, want in test_cases: c = client.Client( tpu='grpc://1.2.3.4:8470', service=self.mock_service_client(tpu_map=tpu_map)) self.assertEqual(want, c.recoverable())
def testRecoverableOOMDisabled(self): test_cases = [ ({ 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'state': 'READY', 'symptoms': [{ 'createTime': '2000-01-01T00:29:30.123456Z', 'symptomType': 'OUT_OF_MEMORY', 'details': 'The TPU runtime has run OOM at timestamp ' '2020-05-29T04:51:32.038721+00:00', 'workerId': '0' }] } }, True), ] FLAGS.runtime_oom_exit = False for tpu_map, want in test_cases: c = client.Client( tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map)) self.assertEqual(want, c.recoverable()) FLAGS.runtime_oom_exit = True
def testWaitForHealthy(self): time_mock = mock.patch.object(time, 'time', autospec=True).start() time_mock.side_effect = self._mock_time sleep_mock = mock.patch.object(time, 'sleep', autospec=True).start() sleep_mock.side_effect = self._mock_sleep health_timeseries = (['UNHEALTHY_MAINTENANCE'] * 30 + ['TIMEOUT'] * 10 + [None] * 20 + ['HEALTHY'] * 30) tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'ipAddress': '10.1.2.3', 'port': '8470', 'state': 'READY', 'health': health_timeseries, }, } c = client.Client(tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map)) # Doesn't throw RuntimeError as TPU becomes HEALTHY before timeout timeout = 80 interval = 5 return_time = 60 c.wait_for_healthy(timeout_s=timeout, interval=interval) self.assertEqual(time.time(), return_time) self.assertEqual(sleep_mock.call_count, return_time / interval)
def testWaitForHealthyRaisesError(self): time_mock = mock.patch.object(time, 'time', autospec=True).start() time_mock.side_effect = self._mock_time sleep_mock = mock.patch.object(time, 'sleep', autospec=True).start() sleep_mock.side_effect = self._mock_sleep # Mock timeseries where takes longer than timeout. health_timeseries = ['UNHEALTHY_MAINTENANCE'] * 50 + ['TIMEOUT'] * 50 tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'ipAddress': '10.1.2.3', 'port': '8470', 'state': 'READY', 'health': health_timeseries, }, } c = client.Client(tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map)) # Doesn't throw RuntimeError as TPU becomes HEALTHY before timeout with self.assertRaisesRegex( RuntimeError, 'Timed out waiting for TPU .* to become healthy'): c.wait_for_healthy(timeout_s=80, interval=5)
def testConfigureTpuVersion(self, urlopen): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'state': 'READY', 'networkEndpoints': [ { 'ipAddress': '1.2.3.4' }, { 'ipAddress': '5.6.7.8' }, ] } } c = client.Client(tpu='tpu_name', project='test-project', zone='us-central1-c', service=self.mock_service_client(tpu_map=tpu_map)) c.configure_tpu_version('1.15') paths = [call[0][0].full_url for call in urlopen.call_args_list] self.assertEqual([ 'http://1.2.3.4:8475/requestversion/1.15', 'http://5.6.7.8:8475/requestversion/1.15' ], sorted(paths))
def testInitializeWithoutMetadata(self): c = client.Client(tpu='tpu_name', project='project', zone='zone') self.assertEqual('tpu_name', c._tpu) self.assertEqual(True, c._use_api) self.assertIsNone(c._service) self.assertIsNone(c._credentials) self.assertEqual('project', c._project) self.assertEqual('zone', c._zone) self.assertIsNone(c._discovery_url)
def testRecoverableNoState(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'ipAddress': '10.1.2.3', 'port': '8470', } } c = client.Client(tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map)) self.assertEqual(True, c.recoverable())
def testInitializeNoArgumentsWithEnvironmentVariable(self): os.environ['TPU_NAME'] = 'tpu_name' tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'ipAddress': '10.1.2.3', 'port': '8470', 'health': 'HEALTHY' } } c = client.Client(service=self.mock_service_client(tpu_map=tpu_map)) self.assertClientContains(c)
def testInitializeTpuName(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'ipAddress': '10.1.2.3', 'port': '8470', 'health': 'HEALTHY' } } c = client.Client(tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map)) self.assertClientContains(c)
def testNetworkEndpointsNotReadyWithApi(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'ipAddress': '10.1.2.3', 'port': '8470', } } c = client.Client(tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map)) self.assertRaisesRegex(RuntimeError, 'TPU .* is not yet ready; state: "None"', c.network_endpoints)
def testInitializeIpAddress(self): c = client.Client(tpu='grpc://1.2.3.4:8470') self.assertEqual('grpc://1.2.3.4:8470', c._tpu) self.assertEqual(False, c._use_api) self.assertIsNone(c._service) self.assertIsNone(c._credentials) self.assertIsNone(c._project) self.assertIsNone(c._zone) self.assertIsNone(c._discovery_url) self.assertEqual([{ 'ipAddress': '1.2.3.4', 'port': '8470' }], c.network_endpoints())
def testAcceleratorTypeApi(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'ipAddress': '10.1.2.3', 'port': '8470', 'state': 'PREEMPTED', 'health': 'HEALTHY', 'acceleratorType': 'v3-8', 'tensorflowVersion': 'nightly', } } c = client.Client(tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map)) self.assertEqual('v3-8', c.accelerator_type())
def testGetTpuVersion(self, urlopen): c = client.Client(tpu='grpc://1.2.3.4:8470') resp = mock.Mock() resp.read.side_effect = ['{}', '{"currentVersion": "someVersion"}'] urlopen.return_value = resp self.assertIsNone(c.runtime_version(), 'Missing key should be handled.') self.assertEqual('someVersion', c.runtime_version(), 'Should return configured version.') paths = [call[0][0].full_url for call in urlopen.call_args_list] self.assertCountEqual([ 'http://1.2.3.4:8475/requestversion', 'http://1.2.3.4:8475/requestversion', ], sorted(paths))
def testInitializeNoArgumentsWithTPUEnvironmentVariableTPUConfig(self): os.environ['TPU_CONFIG'] = json.dumps({ 'project': 'test-project', 'zone': 'us-central1-c', 'tpu_node_name': 'tpu_name', }) tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'ipAddress': '10.1.2.3', 'port': '8470', 'state': 'READY', 'health': 'HEALTHY', } } c = client.Client(service=self.mock_service_client(tpu_map=tpu_map)) self.assertClientContains(c)
def baseConfigureTpuVersion(self): tpu_map = { 'projects/test-project/locations/us-central1-c/nodes/tpu_name': { 'state': 'READY', 'networkEndpoints': [ { 'ipAddress': '1.2.3.4' }, { 'ipAddress': '5.6.7.8' }, ] } } return client.Client(tpu='tpu_name', project='test-project', zone='us-central1-c', service=self.mock_service_client(tpu_map=tpu_map))
def testRecoverableNoApiAccess(self): c = client.Client(tpu='grpc://1.2.3.4:8470') self.assertEqual(True, c.recoverable())
def testInitializeNoArguments(self): with self.assertRaisesRegex( ValueError, 'Please provide a TPU Name to connect to.'): client.Client()
def testInitializeMultiElementTpuArray(self): with self.assertRaisesRegex( NotImplementedError, 'Using multiple TPUs in a single session is not yet implemented' ): client.Client(tpu=['multiple', 'elements'])