def testMixedModeNonOverlappingKey(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"worker": ["worker4:2222", "worker5:2222"]})
cluster_spec_2 = server_lib.ClusterSpec({
"worker": {
3: "worker0:2222",
6: "worker1:2222",
7: "worker2:2222"
}
})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
cluster_spec = union_cluster.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: 'worker4:2222' }
tasks { key: 1 value: 'worker5:2222' }
tasks { key: 3 value: 'worker0:2222' }
tasks { key: 6 value: 'worker1:2222' }
tasks { key: 7 value: 'worker2:2222' }}
"""
self._verifyClusterSpecEquality(cluster_spec, expected_proto)
def testSingleClusterResolver(self):
base_cluster_spec = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
simple_resolver = SimpleClusterResolver(base_cluster_spec)
union_resolver = UnionClusterResolver(simple_resolver)
expected_proto = """
job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
tasks { key: 1 value: 'ps1:2222' } }
job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
tasks { key: 1 value: 'worker1:2222' }
tasks { key: 2 value: 'worker2:2222' } }
"""
actual_cluster_spec = union_resolver.cluster_spec()
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
def testRetainSparseJobWithNoMerging(self):
base_cluster_spec = server_lib.ClusterSpec({
"worker": {
1: "worker0:2222",
3: "worker1:2222",
5: "worker2:2222"
}
})
base_cluster_resolver = SimpleClusterResolver(base_cluster_spec)
union_cluster = UnionClusterResolver(base_cluster_resolver)
cluster_spec = union_cluster.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 1 value: 'worker0:2222' }
tasks { key: 3 value: 'worker1:2222' }
tasks { key: 5 value: 'worker2:2222' } }
"""
self._verifyClusterSpecEquality(cluster_spec, expected_proto)
def testMergedClusterResolverMaster(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"ps": ["ps0:2222", "ps1:2222"]})
cluster_spec_2 = server_lib.ClusterSpec(
{"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
unspecified_master = union_cluster.master()
self.assertEqual(unspecified_master, "")
specified_master = union_cluster.master("worker", 1)
self.assertEqual(specified_master, "worker1:2222")
rpc_master = union_cluster.master("worker", 1, rpc_layer="grpc")
self.assertEqual(rpc_master, "grpc://worker1:2222")
def testOverlappingJobMergedClusterResolver(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"worker": ["worker4:2222", "worker5:2222"]})
cluster_spec_2 = server_lib.ClusterSpec(
{"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
cluster_spec = union_cluster.cluster_spec()
expected_proto = """
job { name: 'worker' tasks { key: 0 value: 'worker4:2222' }
tasks { key: 1 value: 'worker5:2222' }
tasks { key: 2 value: 'worker0:2222' }
tasks { key: 3 value: 'worker1:2222' }
tasks { key: 4 value: 'worker2:2222' } }
"""
self._verifyClusterSpecEquality(cluster_spec, expected_proto)
def testOverlappingDictAndListThrowError(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"worker": ["worker4:2222", "worker5:2222"]})
cluster_spec_2 = server_lib.ClusterSpec({
"worker": {
1: "worker0:2222",
2: "worker1:2222",
3: "worker2:2222"
}
})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
self.assertRaises(KeyError, union_cluster.cluster_spec)
def testOverlappingSparseJobMergedClusterResolverThrowError(self):
cluster_spec_1 = server_lib.ClusterSpec(
{"worker": {
7: "worker4:2222",
9: "worker5:2222"
}})
cluster_spec_2 = server_lib.ClusterSpec({
"worker": {
3: "worker0:2222",
6: "worker1:2222",
7: "worker2:2222"
}
})
cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1)
cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2)
union_cluster = UnionClusterResolver(cluster_resolver_1,
cluster_resolver_2)
self.assertRaises(KeyError, union_cluster.cluster_spec)
def testUnionClusterResolverGetProperties(self):
cluster_spec_1 = server_lib.ClusterSpec({
"ps": ["ps0:2222", "ps1:2222"],
"worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
})
resolver1 = SimpleClusterResolver(cluster_spec_1,
task_type="ps",
task_id=1,
environment="cloud",
num_accelerators={"GPU": 8},
rpc_layer="grpc")
cluster_spec_2 = server_lib.ClusterSpec({
"ps": ["ps2:2222", "ps3:2222"],
"worker": ["worker3:2222", "worker4:2222", "worker5:2222"]
})
resolver2 = SimpleClusterResolver(cluster_spec_2,
task_type="worker",
task_id=2,
environment="local",
num_accelerators={"GPU": 16},
rpc_layer="http")
union_resolver = UnionClusterResolver(resolver1, resolver2)
self.assertEqual(union_resolver.task_type, "ps")
self.assertEqual(union_resolver.task_id, 1)
self.assertEqual(union_resolver.environment, "cloud")
self.assertEqual(union_resolver.num_accelerators(), {"GPU": 8})
self.assertEqual(union_resolver.rpc_layer, "grpc")
union_resolver.task_type = "worker"
union_resolver.task_id = 2
union_resolver.rpc_layer = "http"
self.assertEqual(union_resolver.task_type, "worker")
self.assertEqual(union_resolver.task_id, 2)
self.assertEqual(union_resolver.rpc_layer, "http")
def testUnionMultipleInstanceRetrieval(self):
worker1_name_to_ip = [
{
'name': 'instance1',
'ip': '10.1.2.3'
},
{
'name': 'instance2',
'ip': '10.2.3.4'
},
{
'name': 'instance3',
'ip': '10.3.4.5'
},
]
worker2_name_to_ip = [
{
'name': 'instance4',
'ip': '10.4.5.6'
},
{
'name': 'instance5',
'ip': '10.5.6.7'
},
{
'name': 'instance6',
'ip': '10.6.7.8'
},
]
ps_name_to_ip = [
{
'name': 'ps1',
'ip': '10.100.1.2'
},
{
'name': 'ps2',
'ip': '10.100.2.3'
},
]
worker1_gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='worker',
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(worker1_name_to_ip))
worker2_gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='worker',
port=8470,
credentials=None,
service=self.gen_standard_mock_service_client(worker2_name_to_ip))
ps_gce_cluster_resolver = GCEClusterResolver(
project='test-project',
zone='us-east1-d',
instance_group='test-instance-group',
task_type='ps',
port=2222,
credentials=None,
service=self.gen_standard_mock_service_client(ps_name_to_ip))
union_cluster_resolver = UnionClusterResolver(
worker1_gce_cluster_resolver, worker2_gce_cluster_resolver,
ps_gce_cluster_resolver)
actual_cluster_spec = union_cluster_resolver.cluster_spec()
expected_proto = """
job { name: 'ps' tasks { key: 0 value: '10.100.1.2:2222' }
tasks { key: 1 value: '10.100.2.3:2222' } }
job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' }
tasks { key: 1 value: '10.2.3.4:8470' }
tasks { key: 2 value: '10.3.4.5:8470' }
tasks { key: 3 value: '10.4.5.6:8470' }
tasks { key: 4 value: '10.5.6.7:8470' }
tasks { key: 5 value: '10.6.7.8:8470' } }
"""
self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)