def testGetMasterNoEntries(self):
    tpu_map = {}

    tpu_cluster_resolver = TPUClusterResolver(
        project='test-project',
        zone='us-central1-c',
        tpu_names=[],
        credentials=None,
        service=self.mock_service_client(tpu_map=tpu_map))
    with self.assertRaises(ValueError):
      tpu_cluster_resolver.get_master()
    def testGetMasterNoEntries(self):
        tpu_map = {}

        tpu_cluster_resolver = TPUClusterResolver(
            project='test-project',
            zone='us-central1-c',
            tpu_names=[],
            credentials=None,
            service=self.mock_service_client(tpu_map=tpu_map))
        with self.assertRaises(ValueError):
            tpu_cluster_resolver.get_master()
示例#3
0
 def testGkeEnvironment(self):
     os.environ[
         'KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = 'grpc://10.120.27.5:8470'
     self.assertTrue('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS' in os.environ)
     tpu_cluster_resolver = TPUClusterResolver()
     self.assertTrue(tpu_cluster_resolver._inGke())
     self.assertEqual(compat.as_bytes('grpc://10.120.27.5:8470'),
                      compat.as_bytes(tpu_cluster_resolver._gkeMaster()))
     self.assertEqual(compat.as_bytes('grpc://10.120.27.5:8470'),
                      compat.as_bytes(tpu_cluster_resolver.get_master()))
     del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS']
 def testGkeEnvironment(self):
   os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = 'grpc://10.120.27.5:8470'
   self.assertTrue('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS' in os.environ)
   tpu_cluster_resolver = TPUClusterResolver()
   self.assertTrue(tpu_cluster_resolver._inGke())
   self.assertEqual(
       compat.as_bytes('grpc://10.120.27.5:8470'),
       compat.as_bytes(tpu_cluster_resolver._gkeMaster()))
   self.assertEqual(
       compat.as_bytes('grpc://10.120.27.5:8470'),
       compat.as_bytes(tpu_cluster_resolver.get_master()))
   del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS']
  def testGetMasterMultipleEntries(self):
    tpu_map = {
        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
            'ipAddress': '10.1.2.3',
            'port': '8470'
        },
        'projects/test-project/locations/us-central1-c/nodes/test-tpu-2': {
            'ipAddress': '10.4.5.6',
            'port': '8470'
        }
    }

    tpu_cluster_resolver = TPUClusterResolver(
        project='test-project',
        zone='us-central1-c',
        tpu_names=['test-tpu-2', 'test-tpu-1'],
        credentials=None,
        service=self.mock_service_client(tpu_map=tpu_map))
    self.assertEqual('grpc://10.4.5.6:8470', tpu_cluster_resolver.get_master())
    def testGetMasterMultipleEntries(self):
        tpu_map = {
            'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
                'ipAddress': '10.1.2.3',
                'port': '8470'
            },
            'projects/test-project/locations/us-central1-c/nodes/test-tpu-2': {
                'ipAddress': '10.4.5.6',
                'port': '8470'
            }
        }

        tpu_cluster_resolver = TPUClusterResolver(
            project='test-project',
            zone='us-central1-c',
            tpu_names=['test-tpu-2', 'test-tpu-1'],
            credentials=None,
            service=self.mock_service_client(tpu_map=tpu_map))
        self.assertEqual('grpc://10.4.5.6:8470',
                         tpu_cluster_resolver.get_master())